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>

api/next/62020.txt

@@ -0,0 +1,4 @@
+pkg testing, type TB interface, SetGOMAXPROCS(int) #62020
+pkg testing, method (*T) SetGOMAXPROCS(int) #62020
+pkg testing, method (*B) SetGOMAXPROCS(int) #62020
+pkg testing, method (*F) SetGOMAXPROCS(int) #62020

doc/next/6-stdlib/99-minor/testing/62020.md

@@ -0,0 +1 @@
+The [`SetGOMAXPROCS`](/pkg/testing#T.SetGOMAXPROCS) method changes GOMAXPROCS for the duration of a single test.

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 {

src/testing/testing_test.go

@@ -13,6 +13,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"slices"
 	"strings"
 	"sync"
@@ -258,6 +259,102 @@ func TestSetenvWithParallelGrandParentBeforeSetenv(t *testing.T) {
 	})
 }
 
+func TestSetGOMAXPROCS(t *testing.T) {
+	if runtime.GOARCH == "wasm" {
+		t.Skip("not supported on wasm yet")
+	}
+	tests := []struct {
+		name string
+		newP int
+	}{
+		{
+			name: "overriding value",
+			newP: 1,
+		},
+	}
+
+	for _, test := range tests {
+		p := runtime.GOMAXPROCS(0)
+		t.Run(test.name, func(t *testing.T) {
+			t.SetGOMAXPROCS(test.newP + 1)
+			if runtime.GOMAXPROCS(0) != test.newP+1 {
+				t.Fatalf("unexpected value after t.SetGOMAXPROCS: got %d, want %d", runtime.GOMAXPROCS(0), test.newP+1)
+			}
+		})
+		if runtime.GOMAXPROCS(0) != p {
+			t.Fatalf("unexpected value after t.SetGOMAXPROCS cleanup: got %d, want %d", runtime.GOMAXPROCS(0), p)
+		}
+	}
+}
+
+func TestSetGOMAXPROCSWithParallelAfterSetGOMAXPROCS(t *testing.T) {
+	if runtime.GOARCH == "wasm" {
+		t.Skip("not supported on wasm yet")
+	}
+	defer func() {
+		want := "testing: t.Parallel called after t.SetGOMAXPROCS; cannot set GOMAXPROCS in parallel tests"
+		if got := recover(); got != want {
+			t.Fatalf("expected panic; got %#v want %q", got, want)
+		}
+	}()
+	p := runtime.GOMAXPROCS(0)
+	t.SetGOMAXPROCS(p + 1)
+	t.Parallel()
+}
+
+func TestSetGOMAXPROCSWithParallelBeforeSetGOMAXPROCS(t *testing.T) {
+	if runtime.GOARCH == "wasm" {
+		t.Skip("not supported on wasm yet")
+	}
+	defer func() {
+		want := "testing: t.SetGOMAXPROCS called after t.Parallel; cannot set GOMAXPROCS in parallel tests"
+		if got := recover(); got != want {
+			t.Fatalf("expected panic; got %#v want %q", got, want)
+		}
+	}()
+	t.Parallel()
+	p := runtime.GOMAXPROCS(0)
+	t.SetGOMAXPROCS(p + 1)
+}
+
+func TestSetGOMAXPROCSWithParallelParentBeforeSetGOMAXPROCS(t *testing.T) {
+	if runtime.GOARCH == "wasm" {
+		t.Skip("not supported on wasm yet")
+	}
+	t.Parallel()
+	t.Run("child", func(t *testing.T) {
+		defer func() {
+			want := "testing: t.SetGOMAXPROCS called after t.Parallel; cannot set GOMAXPROCS in parallel tests"
+			if got := recover(); got != want {
+				t.Fatalf("expected panic; got %#v want %q", got, want)
+			}
+		}()
+
+		p := runtime.GOMAXPROCS(0)
+		t.SetGOMAXPROCS(p + 1)
+	})
+}
+
+func TestSetGOMAXPROCSWithParallelGrandParentBeforeSetGOMAXPROCS(t *testing.T) {
+	if runtime.GOARCH == "wasm" {
+		t.Skip("not supported on wasm yet")
+	}
+	t.Parallel()
+	t.Run("child", func(t *testing.T) {
+		t.Run("grand-child", func(t *testing.T) {
+			defer func() {
+				want := "testing: t.SetGOMAXPROCS called after t.Parallel; cannot set GOMAXPROCS in parallel tests"
+				if got := recover(); got != want {
+					t.Fatalf("expected panic; got %#v want %q", got, want)
+				}
+			}()
+
+			p := runtime.GOMAXPROCS(0)
+			t.SetGOMAXPROCS(p + 1)
+		})
+	})
+}
+
 // testingTrueInInit is part of TestTesting.
 var testingTrueInInit = false