[DNM]: for test

.github/workflows/integration.yml

@@ -2,9 +2,9 @@ name: Integration Test
 
 on:
   push:
-    branches: [ master, tidb-6.5 ]
+    branches: [ master, tidb-6.5, tidb-6.5-with-kv-timeout-feature ]
   pull_request:
-    branches: [ master, tidb-6.5 ]
+    branches: [ master, tidb-6.5, tidb-6.5-with-kv-timeout-feature ]
 
 jobs:
 

.github/workflows/test.yml

@@ -2,9 +2,9 @@ name: Unit Test
 
 on:
   push:
-    branches: [ master, tidb-6.5 ]
+    branches: [ master, tidb-6.5, tidb-6.5-with-kv-timeout-feature ]
   pull_request:
-    branches: [ master, tidb-6.5 ]
+    branches: [ master, tidb-6.5, tidb-6.5-with-kv-timeout-feature ]
 
 jobs:
   test:

integration_tests/main_test.go

@@ -27,6 +27,7 @@ func TestMain(m *testing.M) {
 	opts := []goleak.Option{
 		goleak.IgnoreTopFunction("github.com/pingcap/goleveldb/leveldb.(*DB).mpoolDrain"),
 		goleak.IgnoreTopFunction("github.com/tikv/client-go/v2/txnkv/transaction.keepAlive"), // TODO: fix ttlManager goroutine leak
+		goleak.IgnoreTopFunction("github.com/tikv/client-go/v2/config/retry.(*Config).createBackoffFn.newBackoffFn.func2"),
 	}
 
 	goleak.VerifyTestMain(m, opts...)

integration_tests/snapshot_test.go

@@ -287,12 +287,13 @@ func (s *testSnapshotSuite) TestSnapshotThreadSafe() {
 
 func (s *testSnapshotSuite) TestSnapshotRuntimeStats() {
 	reqStats := tikv.NewRegionRequestRuntimeStats()
-	tikv.RecordRegionRequestRuntimeStats(reqStats.Stats, tikvrpc.CmdGet, time.Second)
-	tikv.RecordRegionRequestRuntimeStats(reqStats.Stats, tikvrpc.CmdGet, time.Millisecond)
+	reqStats.RecordRPCRuntimeStats(tikvrpc.CmdGet, time.Second)
+	reqStats.RecordRPCRuntimeStats(tikvrpc.CmdGet, time.Millisecond)
 	snapshot := s.store.GetSnapshot(0)
-	snapshot.SetRuntimeStats(&txnkv.SnapshotRuntimeStats{})
-	snapshot.MergeRegionRequestStats(reqStats.Stats)
-	snapshot.MergeRegionRequestStats(reqStats.Stats)
+	runtimeStats := &txnkv.SnapshotRuntimeStats{}
+	snapshot.SetRuntimeStats(runtimeStats)
+	snapshot.MergeRegionRequestStats(reqStats)
+	snapshot.MergeRegionRequestStats(reqStats)
 	bo := tikv.NewBackofferWithVars(context.Background(), 2000, nil)
 	err := bo.BackoffWithMaxSleepTxnLockFast(5, errors.New("test"))
 	s.Nil(err)

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)
 		}
 	}()
@@ -347,6 +351,12 @@ func (a *batchConn) batchSendLoop(cfg config.TiKVClient) {
 }
 
 func (a *batchConn) getClientAndSend() {
+	if val, err := util.EvalFailpoint("mockBatchClientSendDelay"); err == nil {
+		if timeout, ok := val.(int); ok && timeout > 0 {
+			time.Sleep(time.Duration(timeout * int(time.Millisecond)))
+		}
+	}
+
 	// Choose a connection by round-robbin.
 	var (
 		cli    *batchCommandsClient
@@ -430,7 +440,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")
 		}
@@ -598,7 +608,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)
@@ -781,7 +791,7 @@ func sendBatchRequest(
 	select {
 	case batchConn.batchCommandsCh <- entry:
 	case <-ctx.Done():
-		logutil.BgLogger().Debug("send request is cancelled",
+		logutil.Logger(ctx).Debug("send request is cancelled",
 			zap.String("to", addr), zap.String("cause", ctx.Err().Error()))
 		return nil, errors.WithStack(ctx.Err())
 	case <-timer.C:
@@ -797,7 +807,7 @@ func sendBatchRequest(
 		return tikvrpc.FromBatchCommandsResponse(res)
 	case <-ctx.Done():
 		atomic.StoreInt32(&entry.canceled, 1)
-		logutil.BgLogger().Debug("wait response is cancelled",
+		logutil.Logger(ctx).Debug("wait response is cancelled",
 			zap.String("to", addr), zap.String("cause", ctx.Err().Error()))
 		return nil, errors.WithStack(ctx.Err())
 	case <-timer.C:

internal/client/client_fail_test.go

@@ -36,7 +36,6 @@ package client
 
 import (
 	"context"
-	"fmt"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -47,18 +46,19 @@ 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"
 )
 
 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()
 	rpcClient.option.dialTimeout = time.Second / 3
 
@@ -82,10 +82,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.

internal/client/client_test.go

@@ -55,6 +55,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/internal/logutil"
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"go.uber.org/zap"
@@ -116,11 +117,11 @@ 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()
-	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
+	addr := server.Addr()
 	conn, err := rpcClient.getConnArray(addr, true)
 	assert.Nil(t, err)
 
@@ -240,7 +241,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)
@@ -255,7 +256,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)
@@ -283,7 +284,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)
@@ -308,10 +309,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.
@@ -324,7 +325,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 == "" {
@@ -641,10 +642,10 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
 		conf.TiKVClient.MaxBatchSize = 128
 	})()
 
-	server, port := startMockTikvService()
+	server, port := mock_server.StartMockTikvService()
 	require.True(t, port > 0)
 	require.True(t, server.IsRunning())
-	addr := server.addr
+	addr := server.Addr()
 	client := NewRPCClient()
 	defer func() {
 		err := client.Close()
@@ -681,7 +682,7 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
 	logutil.BgLogger().Info("restart mock tikv server")
 	server.Start(addr)
 	require.True(t, server.IsRunning())
-	require.Equal(t, addr, server.addr)
+	require.Equal(t, addr, server.Addr())
 
 	// Wait batch client to auto reconnect.
 	start := time.Now()
@@ -700,7 +701,7 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
 			cli.unlockForSend()
 			break
 		}
-		if time.Since(start) > time.Second*5 {
+		if time.Since(start) > time.Second*10 {
 			// It shouldn't take too long for batch_client to reconnect.
 			require.Fail(t, "wait batch client reconnect timeout")
 		}

internal/client/main_test.go

@@ -26,6 +26,7 @@ func TestMain(m *testing.M) {
 	opts := []goleak.Option{
 		goleak.IgnoreTopFunction("google.golang.org/grpc.(*ClientConn).WaitForStateChange"),
 		goleak.IgnoreTopFunction("github.com/tikv/client-go/v2/internal/retry.newBackoffFn.func1"),
+		goleak.IgnoreTopFunction("github.com/tikv/client-go/v2/internal/retry.(*Config).createBackoffFn.newBackoffFn.func2"),
 	}
 	goleak.VerifyTestMain(m, opts...)
 }

internal/client/mock_tikv_service_test.go → internal/client/mock_server/mock_tikv_service.go

@@ -18,7 +18,7 @@
 // https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/client/mock_tikv_service_test.go
 //
 
-package client
+package mock_server
 
 import (
 	"context"
@@ -36,7 +36,7 @@ import (
 	"google.golang.org/grpc"
 )
 
-type server struct {
+type MockServer struct {
 	tikvpb.TikvServer
 	grpcServer *grpc.Server
 	addr       string
@@ -49,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
 	}
@@ -94,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 {
@@ -109,16 +116,20 @@ func (s *server) checkMetadata(ctx context.Context) error {
 	return nil
 }
 
-func (s *server) IsRunning() bool {
+func (s *MockServer) IsRunning() bool {
 	return atomic.LoadInt64(&s.running) == 1
 }
 
-func (s *server) Stop() {
+func (s *MockServer) Addr() string {
+	return s.addr
+}
+
+func (s *MockServer) Stop() {
 	s.grpcServer.Stop()
 	atomic.StoreInt64(&s.running, 0)
 }
 
-func (s *server) Start(addr string) int {
+func (s *MockServer) Start(addr string) int {
 	if addr == "" {
 		addr = fmt.Sprintf("%s:%d", "127.0.0.1", 0)
 	}
@@ -148,9 +159,9 @@ func (s *server) Start(addr string) int {
 	return port
 }
 
-// Try to start a gRPC server and retrun the server instance and binded port.
-func startMockTikvService() (*server, int) {
-	server := &server{}
+// 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
 }