refactor

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

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,

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,23 @@ func InitTempFileLogger(level string) (fname string) {
 	log.ReplaceGlobals(lg, p)
 	return fname
 }
+
+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()
+}

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)