ci: Support ut run-multi for CSE

tools/check/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
     importpath = "github.com/pingcap/tidb/tools/check",
     visibility = ["//visibility:private"],
     deps = [
+        "@org_golang_x_sync//errgroup",
         "@org_golang_x_tools//cover",
         "@org_uber_go_automaxprocs//:automaxprocs",
     ],

tools/check/ut.go

@@ -37,6 +37,7 @@ import (
 
 	// Set the correct value when it runs inside docker.
 	_ "go.uber.org/automaxprocs"
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/tools/cover"
 )
 
@@ -68,6 +69,9 @@ ut run $package $test
 // run test cases that match a pattern
 ut run $package 'r:$regex'
 
+// run test cases of multiple packages
+ut run-multi $package1 $package2 ...
+
 // build all test package
 ut build
 
@@ -194,6 +198,30 @@ func cmdBuild(args ...string) bool {
 	return true
 }
 
+func cmdRunMulti(pkgs ...string) bool {
+	var err error
+	if len(pkgs) == 0 {
+		return true
+	}
+
+	// Build tasks
+	var tasks []task
+	start := time.Now()
+	err = buildTestBinaryMulti(pkgs)
+	if err != nil {
+		log.Println("build package error", pkgs, err)
+		return false
+	}
+
+	if tasks, err = listTestCasesForPkgs(pkgs); err != nil {
+		log.Println("run existing test cases error", err)
+		return false
+	}
+
+	fmt.Printf("building task finish, maxproc=%d, count=%d, takes=%v\n", buildParallel, len(tasks), time.Since(start))
+	return runTestCases(tasks)
+}
+
 func cmdRun(args ...string) bool {
 	var err error
 	pkgs, err := listPackages()
@@ -221,7 +249,7 @@ func cmdRun(args ...string) bool {
 				tasks = listLongTasks(pkg, tasks)
 			}
 		} else {
-			if tasks, err = runExistingTestCases(pkgs); err != nil {
+			if tasks, err = listTestCasesForPkgs(pkgs); err != nil {
 				log.Println("run existing test cases error", err)
 				return false
 			}
@@ -319,7 +347,10 @@ func cmdRun(args ...string) bool {
 	}
 
 	fmt.Printf("building task finish, parallelism=%d, count=%d, takes=%v\n", buildParallel, len(tasks), time.Since(start))
+	return runTestCases(tasks)
+}
 
+func runTestCases(tasks []task) bool {
 	testWorkerCount := p
 	if long {
 		testWorkerCount = longTestWorkerCount
@@ -334,7 +365,7 @@ func cmdRun(args ...string) bool {
 
 	shuffle(tasks)
 
-	start = time.Now()
+	start := time.Now()
 	for _, task := range tasks {
 		taskCh <- task
 	}
@@ -368,8 +399,8 @@ func cmdRun(args ...string) bool {
 	return true
 }
 
-func runExistingTestCases(pkgs []string) (tasks []task, err error) {
-	wg := &sync.WaitGroup{}
+func listTestCasesForPkgs(pkgs []string) (tasks []task, err error) {
+	g := new(errgroup.Group)
 	tasksChannel := make(chan []task, len(pkgs))
 	for _, pkg := range pkgs {
 		exist, err := testBinaryExist(pkg)
@@ -381,12 +412,22 @@ func runExistingTestCases(pkgs []string) (tasks []task, err error) {
 			fmt.Println("no test case in ", pkg)
 			continue
 		}
+		pkgCopy := pkg
+		g.Go(func() error {
+			tasks, err := listTestCases(pkgCopy, nil)
+			if err != nil {
+				log.Println("list test cases error", pkgCopy, err)
+				return withTrace(err)
+			}
+			tasksChannel <- tasks
+			return nil
+		})
+	}
 
-		wg.Add(1)
-		go listTestCasesConcurrent(wg, pkg, tasksChannel)
+	if err := g.Wait(); err != nil {
+		return nil, withTrace(err)
 	}
 
-	wg.Wait()
 	close(tasksChannel)
 	for t := range tasksChannel {
 		tasks = append(tasks, t...)
@@ -518,6 +559,8 @@ func main() {
 			isSucceed = cmdBuild(os.Args[2:]...)
 		case "run":
 			isSucceed = cmdRun(os.Args[2:]...)
+		case "run-multi":
+			isSucceed = cmdRunMulti(os.Args[2:]...)
 		default:
 			isSucceed = usage()
 		}
@@ -678,6 +721,7 @@ func (b blocksByStart) Less(i, j int) bool {
 	return bi.StartLine < bj.StartLine || bi.StartLine == bj.StartLine && bi.StartCol < bj.StartCol
 }
 
+// listTestCases list all test cases of a package and append to a slice.
 func listTestCases(pkg string, tasks []task) ([]task, error) {
 	newCases, err := listNewTestCases(pkg)
 	if err != nil {
@@ -691,20 +735,6 @@ func listTestCases(pkg string, tasks []task) ([]task, error) {
 	return tasks, nil
 }
 
-func listTestCasesConcurrent(wg *sync.WaitGroup, pkg string, tasksChannel chan<- []task) {
-	defer wg.Done()
-	newCases, err := listNewTestCases(pkg)
-	if err != nil {
-		log.Println("list test case error", pkg, err)
-		return
-	}
-	tasks := make([]task, 0, len(newCases))
-	for _, c := range newCases {
-		tasks = append(tasks, task{pkg, c})
-	}
-	tasksChannel <- tasks
-}
-
 func filterTestCases(tasks []task, arg1 string) ([]task, error) {
 	if strings.HasPrefix(arg1, "r:") {
 		r, err := regexp.Compile(arg1[2:])