Increase GOMAXPROCS during aio.GoQueue usage

PiperOrigin-RevId: 735048540

Author: nixprime

pkg/aio/BUILD

@@ -13,6 +13,7 @@ go_library(
     ],
     visibility = ["//pkg/sentry:internal"],
     deps = [
+        "//pkg/gomaxprocs",
         "//pkg/sync",
         "@org_golang_x_sys//unix:go_default_library",
     ],

pkg/aio/aio.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/gomaxprocs"
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
@@ -113,6 +114,7 @@ func NewGoQueue(cap int) *GoQueue {
 	for range cap {
 		go q.workerMain()
 	}
+	gomaxprocs.Add(cap)
 	return q
 }
 
@@ -153,6 +155,7 @@ func (q *GoQueue) workerMain() {
 func (q *GoQueue) Destroy() {
 	close(q.shutdown)
 	q.workers.Wait()
+	gomaxprocs.Add(-q.Cap())
 }
 
 // Cap implements Queue.Cap.

pkg/gomaxprocs/BUILD

@@ -0,0 +1,35 @@
+load("//pkg/sync/locking:locking.bzl", "declare_mutex")
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+declare_mutex(
+    name = "gomaxprocs_mutex",
+    out = "gomaxprocs_mutex.go",
+    package = "gomaxprocs",
+    prefix = "gomaxprocs",
+)
+
+go_library(
+    name = "gomaxprocs",
+    srcs = [
+        "gomaxprocs.go",
+        "gomaxprocs_mutex.go",
+    ],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/log",
+        "//pkg/sync",
+        "//pkg/sync/locking",
+    ],
+)
+
+go_test(
+    name = "gomaxprocs_test",
+    size = "small",
+    srcs = ["gomaxprocs_test.go"],
+    library = ":gomaxprocs",
+)

pkg/gomaxprocs/gomaxprocs.go

@@ -0,0 +1,82 @@
+// Copyright 2025 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package gomaxprocs synchronizes adjustments to GOMAXPROCS. When this package
+// is active (i.e. after the first call to SetBase), it sets the value of
+// GOMAXPROCS to a "base" value (which should be set by a single goroutine,
+// without races) plus a non-negative "temporary" value (which may be
+// concurrently increased or decreased by multiple goroutines).
+//
+// Note that changing GOMAXPROCS stops the world, so callers should adjust
+// GOMAXPROCS infrequently.
+//
+// TODO: Add gomaxprocs.Get() and check that other gVisor packages don't call
+// runtime.GOMAXPROCS() at all.
+package gomaxprocs
+
+import (
+	"runtime"
+
+	"gvisor.dev/gvisor/pkg/log"
+)
+
+var (
+	mu gomaxprocsMutex
+	// +checklocks:mu
+	base int
+	// +checklocks:mu
+	temp int
+)
+
+// SetBase sets base GOMAXPROCS.
+func SetBase(n int) {
+	if n < 1 {
+		log.Traceback("Invalid base GOMAXPROCS: %d", n)
+		return
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	oldBase := base
+	base = n
+	updateRuntime(oldBase, temp)
+}
+
+// Add adds n temporary GOMAXPROCS. n may be negative; callers should call Add
+// with negative n to remove temporary GOMAXPROCS when they are no longer
+// needed.
+func Add(n int) {
+	mu.Lock()
+	defer mu.Unlock()
+	t := temp + n
+	if t < 0 {
+		log.Traceback("gomaxprocs.Add(%d) would cause temp to become %d", n, t)
+		return
+	}
+	oldTemp := temp
+	temp = t
+	if base != 0 {
+		updateRuntime(base, oldTemp)
+	}
+}
+
+// +checklocks:mu
+func updateRuntime(oldBase, oldTemp int) {
+	n := base + temp
+	log.Debugf("Setting GOMAXPROCS to %d", n)
+	got := runtime.GOMAXPROCS(n)
+	if want := oldBase + oldTemp; oldBase != 0 && got != want {
+		// Something changed GOMAXPROCS outside of our control.
+		log.Warningf("Previous GOMAXPROCS was %d, expected %d = %d + %d", got, want, oldBase, oldTemp)
+	}
+}

pkg/gomaxprocs/gomaxprocs_test.go

@@ -0,0 +1,84 @@
+// Copyright 2025 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package gomaxprocs
+
+import (
+	"runtime"
+	"testing"
+)
+
+// reset cancels the effect of all previous calls to SetBase and Add and sets
+// GOMAXPROCS to the given value.
+func reset(n int) {
+	mu.Lock()
+	defer mu.Unlock()
+	base = 0
+	temp = 0
+	runtime.GOMAXPROCS(n)
+}
+
+func TestBasic(t *testing.T) {
+	init := runtime.GOMAXPROCS(0)
+	defer reset(init)
+
+	firstBase := init + 1
+	SetBase(firstBase)
+	if got, want := runtime.GOMAXPROCS(0), firstBase; got != want {
+		t.Errorf("GOMAXPROCS after first SetBase(%d): got %d, want %d", firstBase, got, want)
+	}
+	Add(1)
+	if got, want := runtime.GOMAXPROCS(0), firstBase+1; got != want {
+		t.Errorf("GOMAXPROCS after first Add(1): got %d, want %d", got, want)
+	}
+	SetBase(firstBase + 1)
+	if got, want := runtime.GOMAXPROCS(0), firstBase+2; got != want {
+		t.Errorf("GOMAXPROCS after SetBase(%d): got %d, want %d", firstBase+1, got, want)
+	}
+	Add(1)
+	if got, want := runtime.GOMAXPROCS(0), firstBase+3; got != want {
+		t.Errorf("GOMAXPROCS after second Add(1): got %d, want %d", got, want)
+	}
+	SetBase(firstBase)
+	if got, want := runtime.GOMAXPROCS(0), firstBase+2; got != want {
+		t.Errorf("GOMAXPROCS after second SetBase(%d): got %d, want %d", firstBase, got, want)
+	}
+	Add(-2)
+	if got, want := runtime.GOMAXPROCS(0), firstBase; got != want {
+		t.Errorf("GOMAXPROCS after Add(-2): got %d, want %d", got, want)
+	}
+}
+
+func TestAddIgnoredUntilSetBase(t *testing.T) {
+	init := runtime.GOMAXPROCS(0)
+	defer reset(init)
+
+	Add(2)
+	if got, want := runtime.GOMAXPROCS(0), init; got != want {
+		t.Errorf("GOMAXPROCS after Add(2): got %d, want %d", got, want)
+	}
+	Add(1)
+	if got, want := runtime.GOMAXPROCS(0), init; got != want {
+		t.Errorf("GOMAXPROCS after Add(1): got %d, want %d", got, want)
+	}
+	newBase := init + 1
+	SetBase(newBase)
+	if got, want := runtime.GOMAXPROCS(0), newBase+3; got != want {
+		t.Errorf("GOMAXPROCS after SetBase(%d): got %d, want %d", newBase, got, want)
+	}
+	Add(-3)
+	if got, want := runtime.GOMAXPROCS(0), newBase; got != want {
+		t.Errorf("GOMAXPROCS after Add(-3): got %d, want %d", got, want)
+	}
+}

runsc/boot/BUILD

@@ -47,6 +47,7 @@ go_library(
         "//pkg/fd",
         "//pkg/flipcall",
         "//pkg/fspath",
+        "//pkg/gomaxprocs",
         "//pkg/hostos",
         "//pkg/log",
         "//pkg/memutil",

runsc/boot/loader.go

@@ -34,6 +34,7 @@ import (
 	"gvisor.dev/gvisor/pkg/coverage"
 	"gvisor.dev/gvisor/pkg/cpuid"
 	"gvisor.dev/gvisor/pkg/fd"
+	"gvisor.dev/gvisor/pkg/gomaxprocs"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/memutil"
 	"gvisor.dev/gvisor/pkg/metric"
@@ -560,7 +561,7 @@ func New(args Args) (*Loader, error) {
 		args.NumCPU = runtime.NumCPU()
 	}
 	log.Infof("CPUs: %d", args.NumCPU)
-	runtime.GOMAXPROCS(args.NumCPU)
+	gomaxprocs.SetBase(args.NumCPU)
 
 	if args.TotalHostMem > 0 {
 		// As per tmpfs(5), the default size limit is 50% of total physical RAM.