testing: add TB.SetGOMAXPROCS function

Add a new method TB.SetGOMAXPROCS which sets variable of GOMAXPROCS.
This method aims to set a variable for the isolated lifetime of the test and cleans up.
And unset this when the test ends.
This method disables the test or benchmark from running in
parallel.

Fixes: #62020

Change-Id: Iae44109d0def35cc47049c3ca4cd5306173d52ee
Signed-off-by: sivchari <shibuuuu5@gmail.com>
Reviewed-on: https://go-review.googlesource.com/c/go/+/519235
Reviewed-by: Bryan Mills <bcmills@google.com>
Reviewed-by: Ian Lance Taylor <iant@google.com>
Auto-Submit: Ian Lance Taylor <iant@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

1 files changed, 47 insertions, 2 deletions

diff --git a/src/testing/testing.go b/src/testing/testing.go
index 5c06aea5f8..a5441f24d5 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -890,6 +890,7 @@ type TB interface {
 	Logf(format string, args ...any)
 	Name() string
 	Setenv(key, value string)
+	SetGOMAXPROCS(n int)
 	Skip(args ...any)
 	SkipNow()
 	Skipf(format string, args ...any)
@@ -916,8 +917,9 @@ var _ TB = (*B)(nil)
 // may be called simultaneously from multiple goroutines.
 type T struct {
 	common
-	isEnvSet bool
-	context  *testContext // For running tests and subtests.
+	isEnvSet        bool
+	isGOMAXPROCSSet bool
+	context         *testContext // For running tests and subtests.
 }
 
 func (c *common) private() {}
@@ -1306,6 +1308,19 @@ func (c *common) Setenv(key, value string) {
 	}
 }
 
+// SetGOMAXPROCS calls runtime.GOMAXPROCS(n) and uses Cleanup to
+// restore the value of GOMAXPROCS after the test.
+//
+// Because GOMAXPROCS affects the whole process, it cannot be used
+// in parallel tests or tests with parallel ancestors.
+func (c *common) SetGOMAXPROCS(n int) {
+	c.checkFuzzFn("SetGOMAXPROCS")
+	prev := runtime.GOMAXPROCS(n)
+	c.Cleanup(func() {
+		runtime.GOMAXPROCS(prev)
+	})
+}
+
 // panicHanding controls the panic handling used by runCleanup.
 type panicHandling int
 
@@ -1446,6 +1461,9 @@ func (t *T) Parallel() {
 	if t.isEnvSet {
 		panic("testing: t.Parallel called after t.Setenv; cannot set environment variables in parallel tests")
 	}
+	if t.isGOMAXPROCSSet {
+		panic("testing: t.Parallel called after t.SetGOMAXPROCS; cannot set GOMAXPROCS in parallel tests")
+	}
 	t.isParallel = true
 	if t.parent.barrier == nil {
 		// T.Parallel has no effect when fuzzing.
@@ -1527,6 +1545,33 @@ func (t *T) Setenv(key, value string) {
 	t.common.Setenv(key, value)
 }
 
+// SetGOMAXPROCS calls runtime.GOMAXPROCS(n) and uses Cleanup to
+// restore the value of GOMAXPROCS after the test.
+//
+// Because GOMAXPROCS affects the whole process, it cannot be used
+// in parallel tests or tests with parallel ancestors.
+func (t *T) SetGOMAXPROCS(n int) {
+	// Non-parallel subtests that have parallel ancestors may still
+	// run in parallel with other tests: they are only non-parallel
+	// with respect to the other subtests of the same parent.
+	// Since SetGOMAXPROCS affects the whole process, we need to disallow it
+	// if the current test or any parent is parallel.
+	isParallel := false
+	for c := &t.common; c != nil; c = c.parent {
+		if c.isParallel {
+			isParallel = true
+			break
+		}
+	}
+	if isParallel {
+		panic("testing: t.SetGOMAXPROCS called after t.Parallel; cannot set GOMAXPROCS in parallel tests")
+	}
+
+	t.isGOMAXPROCSSet = true
+
+	t.common.SetGOMAXPROCS(n)
+}
+
 // InternalTest is an internal type but exported because it is cross-package;
 // it is part of the implementation of the "go test" command.
 type InternalTest struct {