From c540f269ab3375e7722b7e79578eccd62d78fed3 Mon Sep 17 00:00:00 2001
From: wangyueliang <wangyueliang@kylinos.cn>
Date: Wed, 8 May 2024 14:07:52 +0800
Subject: [PATCH] kola support splitting different test cases for concurrent
 execution

[upstream]
5ffbf376b9dfaffb3a71c25b234714fcb5d4d7fc
---
 mantle/cmd/kola/options.go |  1 +
 mantle/harness/suite.go    |  9 ++++++-
 mantle/kola/harness.go     | 48 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/mantle/cmd/kola/options.go b/mantle/cmd/kola/options.go
index bfd86894..e5405062 100644
--- a/mantle/cmd/kola/options.go
+++ b/mantle/cmd/kola/options.go
@@ -63,6 +63,7 @@ func init() {
 	bv(&kola.NoNet, "no-net", false, "Don't run tests that require an Internet connection")
 	bv(&kola.ForceRunPlatformIndependent, "run-platform-independent", false, "Run tests that claim platform independence")
 	ssv(&kola.Tags, "tag", []string{}, "Test tag to run. Can be specified multiple times.")
+	sv(&kola.Sharding, "sharding", "", "Provide e.g. 'hash:m/n' where m and n are integers, 1 <= m <= n.  Only tests hashing to m will be run.")
 	bv(&kola.Options.SSHOnTestFailure, "ssh-on-test-failure", false, "SSH into a machine when tests fail")
 	sv(&kola.Options.Stream, "stream", "", "CoreOS stream ID (e.g. for Fedora CoreOS: stable, testing, next)")
 	sv(&kola.Options.CosaWorkdir, "workdir", "", "coreos-assembler working directory")
diff --git a/mantle/harness/suite.go b/mantle/harness/suite.go
index c4189461..fc144610 100644
--- a/mantle/harness/suite.go
+++ b/mantle/harness/suite.go
@@ -75,6 +75,9 @@ type Options struct {
 	// Limit number of tests to run in parallel (0 means GOMAXPROCS).
 	Parallel int
 
+	// Sharding splits tests across runners
+	Sharding string
+
 	Reporters reporters.Reporters
 }
 
@@ -297,7 +300,11 @@ func (s *Suite) runTests(out, tap io.Writer) error {
 		go func() { <-t.signal }()
 	})
 	if !t.ran {
-		return SuiteEmpty
+		if s.opts.Sharding != "" {
+			fmt.Printf("notice: sharding %s enabled, no tests matched\n", s.opts.Sharding)
+		} else {
+			return SuiteEmpty
+		}
 	}
 	if t.Failed() {
 		s.opts.Reporters.SetResult(testresult.Fail)
diff --git a/mantle/kola/harness.go b/mantle/kola/harness.go
index 176acf35..2a732deb 100644
--- a/mantle/kola/harness.go
+++ b/mantle/kola/harness.go
@@ -131,6 +131,9 @@ var (
 	WarnOnErrorTests    []string // denylisted tests we are going to run and warn in case of error
 	Tags                []string // tags to be ran
 
+	// Sharding is a string of the form: hash:m/n where m and n are integers to run only tests which hash to m.
+	Sharding string
+
 	extTestNum  = 1 // Assigns a unique number to each non-exclusive external test
 	testResults protectedTestResults
 
@@ -796,9 +799,15 @@ func runProvidedTests(testsBank map[string]*register.Test, patterns []string, mu
 		tests = newTests
 	}
 
+	tests, err = shardTests(tests, Sharding)
+	if err != nil {
+		plog.Fatalf("%v", err)
+	}
+
 	opts := harness.Options{
 		OutputDir: outputDir,
 		Parallel:  TestParallelism,
+		Sharding:  Sharding,
 		Verbose:   true,
 		Reporters: reporters.Reporters{
 			reporters.NewJSONReporter("report.json", pltfrm, versionStr),
@@ -1543,6 +1552,45 @@ mainloop:
 	return buckets
 }
 
+// shardTests filters tests to a particular shard - i.e. a group of tests
+// whose name hashes to the same value.
+func shardTests(tests map[string]*register.Test, sharding string) (map[string]*register.Test, error) {
+	if sharding == "" {
+		return tests, nil
+	}
+	if !strings.HasPrefix(sharding, "hash:") {
+		return nil, fmt.Errorf("invalid sharding syntax: %s", sharding)
+	}
+	parts := strings.SplitN(sharding[len("hash:"):], "/", 2)
+	if len(parts) != 2 {
+		return nil, fmt.Errorf("invalid sharding syntax: %s", sharding)
+	}
+	mv, err := strconv.Atoi(parts[0])
+	if err != nil {
+		return nil, fmt.Errorf("invalid sharding syntax '%s': %w", sharding, err)
+	}
+	nv, err := strconv.Atoi(parts[1])
+	if err != nil {
+		return nil, fmt.Errorf("invalid sharding syntax '%s': %w", sharding, err)
+	}
+	if mv > nv || nv < 1 || mv < 1 {
+		return nil, fmt.Errorf("invalid sharding in '%s'", sharding)
+	}
+	m := uint(mv)
+	n := uint(nv)
+
+	ret := make(map[string]*register.Test)
+	for name, test := range tests {
+		h := fnv.New64()
+		h.Write([]byte(name))
+		d := uint(h.Sum64()%uint64(n)) + 1
+		if d == m {
+			ret[name] = test
+		}
+	}
+	return ret, nil
+}
+
 // Create a parent test that runs non-exclusive tests as subtests
 func makeNonExclusiveTest(bucket int, tests []*register.Test, flight platform.Flight) register.Test {
 	// Parse test flags and gather configs
-- 
Gitee