From c4546cb1bfe0fa961ccf58ccc9b0c7131cfa5906 Mon Sep 17 00:00:00 2001
From: crazycs520 <crazycs520@gmail.com>
Date: Wed, 13 Mar 2024 16:54:24 +0800
Subject: [PATCH] refine test

Signed-off-by: crazycs520 <crazycs520@gmail.com>
---
 internal/client/client_batch.go |  1 +
 internal/client/client_test.go  | 11 ++++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/internal/client/client_batch.go b/internal/client/client_batch.go
index bd21690a58..e575f05350 100644
--- a/internal/client/client_batch.go
+++ b/internal/client/client_batch.go
@@ -391,6 +391,7 @@ func (a *batchConn) getClientAndSend() {
 			time.Sleep(time.Duration(timeout * int(time.Millisecond)))
 		}
 	}
+
 	// Choose a connection by round-robbin.
 	var (
 		cli    *batchCommandsClient
diff --git a/internal/client/client_test.go b/internal/client/client_test.go
index 4ab2e90a79..54db795cae 100644
--- a/internal/client/client_test.go
+++ b/internal/client/client_test.go
@@ -59,6 +59,7 @@ import (
 	"github.com/tikv/client-go/v2/internal/client/mockserver"
 	"github.com/tikv/client-go/v2/internal/logutil"
 	"github.com/tikv/client-go/v2/tikvrpc"
+	"github.com/tikv/client-go/v2/util/israce"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/metadata"
@@ -881,6 +882,10 @@ func TestBatchClientReceiveHealthFeedback(t *testing.T) {
 }
 
 func TestConcurrencyRequestLimitWithEnableForwarding(t *testing.T) {
+	if israce.RaceEnabled {
+		// The test will failed when run with race. see https://github.com/tikv/client-go/issues/1222
+		t.Skip()
+	}
 	store1, port1 := mockserver.StartMockTikvService()
 	require.True(t, port1 > 0)
 	require.True(t, store1.IsRunning())
@@ -901,7 +906,7 @@ func TestConcurrencyRequestLimitWithEnableForwarding(t *testing.T) {
 	wg := sync.WaitGroup{}
 	wg.Add(1)
 	done := int64(0)
-	concurrency := 200
+	concurrency := 500
 	go func() {
 		defer wg.Done()
 		for {
@@ -922,7 +927,7 @@ func TestConcurrencyRequestLimitWithEnableForwarding(t *testing.T) {
 				atomic.AddInt64(&done, 1)
 				wg.Done()
 			}()
-			for i := 0; i < 2000; i++ {
+			for i := 0; i < 5000; i++ {
 				req := &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Coprocessor{Coprocessor: &coprocessor.Request{}}}
 				forwardedHost := ""
 				if i%2 != 0 {
@@ -947,6 +952,6 @@ func TestConcurrencyRequestLimitWithEnableForwarding(t *testing.T) {
 		require.Equal(t, int64(9223372036854775807), cli.maxConcurrencyRequestLimit.Load())
 		require.True(t, cli.available() > 0, fmt.Sprintf("sent: %d", cli.sent.Load()))
 		// TODO(crazycs520): fix following assertion. cli.sent may be less than 0 in some cases when enable-forwarding is true.
-		//require.True(t, cli.sent.Load() >= 0, fmt.Sprintf("sent: %d", cli.sent.Load()))
+		require.True(t, cli.sent.Load() >= 0, fmt.Sprintf("sent: %d", cli.sent.Load()))
 	}
 }