From 39108f7be343368ed7487848d81d4528734dc02a Mon Sep 17 00:00:00 2001
From: lhy1024 <admin@liudos.us>
Date: Mon, 25 Mar 2024 15:05:16 +0800
Subject: [PATCH] tests/etcdutil: reduce `TestRandomKillEtcd` test time (#7947)

ref tikv/pd#7930

Signed-off-by: lhy1024 <admin@liudos.us>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
---
 pkg/utils/etcdutil/etcdutil_test.go | 28 ++++++--------------------
 pkg/utils/testutil/testutil.go      | 23 +++++++++++++++++++++
 server/api/region_test.go           | 31 ++++++++---------------------
 3 files changed, 37 insertions(+), 45 deletions(-)

diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go
index 4fb96895942..e02615b695f 100644
--- a/pkg/utils/etcdutil/etcdutil_test.go
+++ b/pkg/utils/etcdutil/etcdutil_test.go
@@ -239,7 +239,7 @@ func TestRandomKillEtcd(t *testing.T) {
 
 	// Randomly kill an etcd server and restart it
 	cfgs := []embed.Config{etcds[0].Config(), etcds[1].Config(), etcds[2].Config()}
-	for i := 0; i < 10; i++ {
+	for i := 0; i < len(cfgs)*2; i++ {
 		killIndex := rand.Intn(len(etcds))
 		etcds[killIndex].Close()
 		checkEtcdEndpointNum(re, client1, 2)
@@ -452,9 +452,9 @@ func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() {
 	re := suite.Require()
 	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/utils/etcdutil/meetEtcdError", `return()`))
 	cache := make(map[string]struct{})
-	for i := 0; i < int(maxLoadBatchSize)*2; i++ {
+	testutil.GenerateTestDataConcurrently(int(maxLoadBatchSize)*2, func(i int) {
 		suite.put(re, fmt.Sprintf("TestLoadWithLimitChange%d", i), "")
-	}
+	})
 	watcher := NewLoopWatcher(
 		suite.ctx,
 		&suite.wg,
@@ -583,25 +583,9 @@ func (suite *loopWatcherTestSuite) TestWatcherLoadLargeKey() {
 	count := 65536
 	ctx, cancel := context.WithCancel(suite.ctx)
 	defer cancel()
-
-	// create data
-	var wg sync.WaitGroup
-	tasks := make(chan int, count)
-	for w := 0; w < 16; w++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			for i := range tasks {
-				suite.put(re, fmt.Sprintf("TestWatcherLoadLargeKey/test-%d", i), "")
-			}
-		}()
-	}
-	for i := 0; i < count; i++ {
-		tasks <- i
-	}
-	close(tasks)
-	wg.Wait()
-
+	testutil.GenerateTestDataConcurrently(count, func(i int) {
+		suite.put(re, fmt.Sprintf("TestWatcherLoadLargeKey/test-%d", i), "")
+	})
 	cache := make([]string, 0)
 	watcher := NewLoopWatcher(
 		ctx,
diff --git a/pkg/utils/testutil/testutil.go b/pkg/utils/testutil/testutil.go
index a41fc436ca6..cef952353bc 100644
--- a/pkg/utils/testutil/testutil.go
+++ b/pkg/utils/testutil/testutil.go
@@ -16,7 +16,9 @@ package testutil
 
 import (
 	"os"
+	"runtime"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/pingcap/kvproto/pkg/pdpb"
@@ -101,3 +103,24 @@ func InitTempFileLogger(level string) (fname string) {
 	log.ReplaceGlobals(lg, p)
 	return fname
 }
+
+// GenerateTestDataConcurrently generates test data concurrently.
+func GenerateTestDataConcurrently(count int, f func(int)) {
+	var wg sync.WaitGroup
+	tasks := make(chan int, count)
+	workers := runtime.NumCPU()
+	for w := 0; w < workers; w++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := range tasks {
+				f(i)
+			}
+		}()
+	}
+	for i := 0; i < count; i++ {
+		tasks <- i
+	}
+	close(tasks)
+	wg.Wait()
+}
diff --git a/server/api/region_test.go b/server/api/region_test.go
index e10bfbd1af0..4198cdcb694 100644
--- a/server/api/region_test.go
+++ b/server/api/region_test.go
@@ -23,7 +23,6 @@ import (
 	"net/http"
 	"net/url"
 	"sort"
-	"sync"
 	"testing"
 	"time"
 
@@ -333,29 +332,15 @@ func TestRegionsWithKillRequest(t *testing.T) {
 	addr := svr.GetAddr()
 	url := fmt.Sprintf("%s%s/api/v1/regions", addr, apiPrefix)
 	mustBootstrapCluster(re, svr)
-	regionCount := 100000
 
-	// create data
-	var wg sync.WaitGroup
-	tasks := make(chan int, regionCount)
-	for w := 0; w < 16; w++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			for i := range tasks {
-				r := core.NewTestRegionInfo(uint64(i+2), 1,
-					[]byte(fmt.Sprintf("%09d", i)),
-					[]byte(fmt.Sprintf("%09d", i+1)),
-					core.SetApproximateKeys(10), core.SetApproximateSize(10))
-				mustRegionHeartbeat(re, svr, r)
-			}
-		}()
-	}
-	for i := 0; i < regionCount; i++ {
-		tasks <- i
-	}
-	close(tasks)
-	wg.Wait()
+	regionCount := 100000
+	tu.GenerateTestDataConcurrently(regionCount, func(i int) {
+		r := core.NewTestRegionInfo(uint64(i+2), 1,
+			[]byte(fmt.Sprintf("%09d", i)),
+			[]byte(fmt.Sprintf("%09d", i+1)),
+			core.SetApproximateKeys(10), core.SetApproximateSize(10))
+		mustRegionHeartbeat(re, svr, r)
+	})
 
 	ctx, cancel := context.WithCancel(context.Background())
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, http.NoBody)