tools/ut: add a parallel parameter (tikv#8186)

ref tikv#7969 add parallel parameter Signed-off-by: husharp <[email protected]> Signed-off-by: okJiang <[email protected]>
okJiang · May 23, 2024 · 3c25cf9 · 3c25cf9
1 parent 32f75b1
commit 3c25cf9
Showing 1 changed file with 23 additions and 12 deletions.
diff --git a/tools/pd-ut/ut.go b/tools/pd-ut/ut.go
@@ -95,8 +95,7 @@ var (
 
 var (
 	// runtime
-	p                int
-	buildParallel    int
+	parallel         int
 	workDir          string
 	coverFileTempDir string
 	// arguments
@@ -108,6 +107,7 @@ var (
 
 func main() {
 	race = handleFlag("--race")
+	parallelStr := stripFlag("--parallel")
 	junitFile = stripFlag("--junitfile")
 	coverProfile = stripFlag("--coverprofile")
 	ignoreDir = stripFlag("--ignore")
@@ -122,11 +122,21 @@ func main() {
 		defer os.RemoveAll(coverFileTempDir)
 	}
 
-	// Get the correct count of CPU if it's in docker.
-	p = runtime.GOMAXPROCS(0)
-	// We use 2 * p for `go build` to make it faster.
-	buildParallel = p * 2
 	var err error
+	procs := runtime.GOMAXPROCS(0)
+	if parallelStr == "" {
+		// Get the correct count of CPU if it's in docker.
+		parallel = procs
+	} else {
+		parallel, err = strconv.Atoi(parallelStr)
+		if err != nil {
+			fmt.Println("parse parallel error", err)
+			return
+		}
+		if parallel > procs {
+			fmt.Printf("Recommend to set parallel be same as the GOMAXPROCS=%d\n", procs)
+		}
+	}
 	workDir, err = os.Getwd()
 	if err != nil {
 		fmt.Println("os.Getwd() error", err)
@@ -353,12 +363,12 @@ func cmdRun(args ...string) bool {
 		}
 	}
 
-	fmt.Printf("building task finish, parallelism=%d, count=%d, takes=%v\n", buildParallel, len(tasks), time.Since(start))
+	fmt.Printf("building task finish, parallelism=%d, count=%d, takes=%v\n", parallel*2, len(tasks), time.Since(start))
 
 	taskCh := make(chan task, 100)
-	works := make([]numa, p)
+	works := make([]numa, parallel)
 	var wg sync.WaitGroup
-	for i := 0; i < p; i++ {
+	for i := 0; i < parallel; i++ {
 		wg.Add(1)
 		go works[i].worker(&wg, taskCh)
 	}
@@ -400,7 +410,7 @@ func cmdRun(args ...string) bool {
 
 // stripFlag strip the '--flag xxx' from the command line os.Args
 // Example of the os.Args changes
-// Before: ut run pkg TestXXX --coverprofile xxx --junitfile yyy
+// Before: ut run pkg TestXXX --coverprofile xxx --junitfile yyy --parallel 16
 // After: ut run pkg TestXXX
 // The value of the flag is returned.
 func stripFlag(flag string) string {
@@ -636,7 +646,7 @@ func (*numa) testCommand(pkg string, fn string) *exec.Cmd {
 		args = append(args, "-test.coverprofile", tmpFile)
 	}
 	if strings.Contains(fn, "Suite") {
-		args = append(args, "-test.cpu", fmt.Sprint(p/2))
+		args = append(args, "-test.cpu", fmt.Sprint(parallel/2))
 	} else {
 		args = append(args, "-test.cpu", "1")
 	}
@@ -705,7 +715,8 @@ func buildTestBinaryMulti(pkgs []string) error {
 		packages = append(packages, path.Join(modulePath, pkg))
 	}
 
-	p := strconv.Itoa(buildParallel)
+	// We use 2 * parallel for `go build` to make it faster.
+	p := strconv.Itoa(parallel * 2)
 	cmd := exec.Command("go", "test", "-p", p, "--exec", xprogPath, "-vet", "off", "--tags=tso_function_test,deadlock")
 	if coverProfile != "" {
 		coverpkg := "./..."