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, 97 insertions, 0 deletions

diff --git a/src/testing/testing_test.go b/src/testing/testing_test.go
index d3822dfd57..28b5809eea 100644
--- 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