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>

diff --git a/api/next/62020.txt b/api/next/62020.txt
new file mode 100644 (file)
index 0000000..3820d88
--- /dev/null
+++ b/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

diff --git a/doc/next/6-stdlib/99-minor/testing/62020.md b/doc/next/6-stdlib/99-minor/testing/62020.md
new file mode 100644 (file)
index 0000000..3a73a1d
--- /dev/null
+++ b/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.

diff --git a/src/testing/testing.go b/src/testing/testing.go
index 5c06aea5f8fe6b8b0d18290a87a61da6d239bd9b..a5441f24d5c23cd53f486dde1cbd159e7a7dffc7 100644 (file)
--- 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 {

diff --git a/src/testing/testing_test.go b/src/testing/testing_test.go
index d3822dfd578f803bf45cb196429f0ff997ce6016..28b5809eea95c30920c7967e7179d5c9f505f671 100644 (file)
--- a/src/testing/testing_test.go
+++ b/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