2 files changed, 215 insertions, 0 deletions
diff --git a/src/runtime/testdata/testprog/gomaxprocs.go b/src/runtime/testdata/testprog/gomaxprocs.go
new file mode 100644
index 0000000000..915e3c4dad
--- /dev/null
+++ b/src/runtime/testdata/testprog/gomaxprocs.go
@@ -0,0 +1,152 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"strconv"
+	"time"
+)
+
+func init() {
+	register("PrintGOMAXPROCS", PrintGOMAXPROCS)
+	register("SetLimitThenDefaultGOMAXPROCS", SetLimitThenDefaultGOMAXPROCS)
+	register("UpdateGOMAXPROCS", UpdateGOMAXPROCS)
+	register("DontUpdateGOMAXPROCS", DontUpdateGOMAXPROCS)
+}
+
+func PrintGOMAXPROCS() {
+	println(runtime.GOMAXPROCS(0))
+}
+
+func mustSetCPUMax(path string, quota int64) {
+	q := "max"
+	if quota >= 0 {
+		q = strconv.FormatInt(quota, 10)
+	}
+	buf := fmt.Sprintf("%s 100000", q)
+	if err := os.WriteFile(path, []byte(buf), 0); err != nil {
+		panic(fmt.Sprintf("error setting cpu.max: %v", err))
+	}
+}
+
+func mustParseInt64(s string) int64 {
+	v, err := strconv.ParseInt(s, 10, 64)
+	if err != nil {
+		panic(err)
+	}
+	return v
+}
+
+// Inputs:
+// GO_TEST_CPU_MAX_PATH: Path to cgroup v2 cpu.max file.
+// GO_TEST_CPU_MAX_QUOTA: CPU quota to set.
+func SetLimitThenDefaultGOMAXPROCS() {
+	path := os.Getenv("GO_TEST_CPU_MAX_PATH")
+	quota := mustParseInt64(os.Getenv("GO_TEST_CPU_MAX_QUOTA"))
+
+	mustSetCPUMax(path, quota)
+
+	runtime.SetDefaultGOMAXPROCS()
+	println(runtime.GOMAXPROCS(0))
+}
+
+// Wait for GOMAXPROCS to change from from to to. Times out after 10s.
+func waitForMaxProcsChange(from, to int) {
+	start := time.Now()
+	for {
+		if time.Since(start) > 10*time.Second {
+			panic("no update for >10s")
+		}
+
+		procs := runtime.GOMAXPROCS(0)
+		println("GOMAXPROCS:", procs)
+		if procs == to {
+			return
+		}
+		if procs != from {
+			panic(fmt.Sprintf("GOMAXPROCS change got %d want %d", procs, to))
+		}
+
+		time.Sleep(100*time.Millisecond)
+	}
+}
+
+// Make sure that GOMAXPROCS does not change from curr.
+//
+// It is impossible to assert that it never changes, so this just makes sure it
+// stays for 5s.
+func mustNotChangeMaxProcs(curr int) {
+	start := time.Now()
+	for {
+		if time.Since(start) > 5*time.Second {
+			return
+		}
+
+		procs := runtime.GOMAXPROCS(0)
+		println("GOMAXPROCS:", procs)
+		if procs != curr {
+			panic(fmt.Sprintf("GOMAXPROCS change got %d want %d", procs, curr))
+		}
+
+		time.Sleep(100*time.Millisecond)
+	}
+}
+
+// Inputs:
+// GO_TEST_CPU_MAX_PATH: Path to cgroup v2 cpu.max file.
+func UpdateGOMAXPROCS() {
+	// We start with no limit.
+
+	ncpu := runtime.NumCPU()
+
+	procs := runtime.GOMAXPROCS(0)
+	println("GOMAXPROCS:", procs)
+	if procs != ncpu {
+		panic(fmt.Sprintf("GOMAXPROCS got %d want %d", procs, ncpu))
+	}
+
+	path := os.Getenv("GO_TEST_CPU_MAX_PATH")
+
+	// Drop down to 3 CPU.
+	mustSetCPUMax(path, 300000)
+	waitForMaxProcsChange(ncpu, 3)
+
+	// Drop even further. Now we hit the minimum GOMAXPROCS=2.
+	mustSetCPUMax(path, 100000)
+	waitForMaxProcsChange(3, 2)
+
+	// Increase back up.
+	mustSetCPUMax(path, 300000)
+	waitForMaxProcsChange(2, 3)
+
+	// Remove limit entirely.
+	mustSetCPUMax(path, -1)
+	waitForMaxProcsChange(3, ncpu)
+
+	// Setting GOMAXPROCS explicitly disables updates.
+	runtime.GOMAXPROCS(3)
+	mustSetCPUMax(path, 200000)
+	mustNotChangeMaxProcs(3)
+
+	println("OK")
+}
+
+// Inputs:
+// GO_TEST_CPU_MAX_PATH: Path to cgroup v2 cpu.max file.
+func DontUpdateGOMAXPROCS() {
+	// The caller has disabled updates. Make sure they don't happen.
+
+	curr := runtime.GOMAXPROCS(0)
+	println("GOMAXPROCS:", curr)
+
+	path := os.Getenv("GO_TEST_CPU_MAX_PATH")
+	mustSetCPUMax(path, 300000)
+	mustNotChangeMaxProcs(curr)
+
+	println("OK")
+}
diff --git a/src/runtime/testdata/testprog/gomaxprocs_windows.go b/src/runtime/testdata/testprog/gomaxprocs_windows.go
new file mode 100644
index 0000000000..bc7a4b1063
--- /dev/null
+++ b/src/runtime/testdata/testprog/gomaxprocs_windows.go
@@ -0,0 +1,63 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"os"
+	"runtime"
+	"syscall"
+	"unsafe"
+)
+
+func init() {
+	register("WindowsUpdateGOMAXPROCS", WindowsUpdateGOMAXPROCS)
+	register("WindowsDontUpdateGOMAXPROCS", WindowsDontUpdateGOMAXPROCS)
+}
+
+// Set CPU affinity mask to only two CPUs.
+//
+// Skips the test if CPUs 0 and 1 are not available.
+func setAffinity2() {
+	kernel32 := syscall.MustLoadDLL("kernel32.dll")
+	_GetProcessAffinityMask := kernel32.MustFindProc("GetProcessAffinityMask")
+	_SetProcessAffinityMask := kernel32.MustFindProc("SetProcessAffinityMask")
+
+	h, err := syscall.GetCurrentProcess()
+	if err != nil {
+		panic(err)
+	}
+
+	var mask, sysmask uintptr
+	ret, _, err := _GetProcessAffinityMask.Call(uintptr(h), uintptr(unsafe.Pointer(&mask)), uintptr(unsafe.Pointer(&sysmask)))
+	if ret == 0 {
+		panic(err)
+	}
+
+	// We're going to restrict to CPUs 0 and 1. Make sure those are already available.
+	if mask & 0b11 != 0b11 {
+		println("SKIP: CPUs 0 and 1 not available")
+		os.Exit(0)
+	}
+
+	mask = 0b11
+	ret, _, err = _SetProcessAffinityMask.Call(uintptr(h), mask)
+	if ret == 0 {
+		panic(err)
+	}
+}
+
+func WindowsUpdateGOMAXPROCS() {
+	ncpu := runtime.NumCPU()
+	setAffinity2()
+	waitForMaxProcsChange(ncpu, 2)
+	println("OK")
+}
+
+func WindowsDontUpdateGOMAXPROCS() {
+	ncpu := runtime.NumCPU()
+	setAffinity2()
+	mustNotChangeMaxProcs(ncpu)
+	println("OK")
+}