From 211bb0f8832582b40cc810af493c4250fa9fd3ab Mon Sep 17 00:00:00 2001
From: Konstantin Bogomolov <bogomolov@google.com>
Date: Mon, 13 May 2024 20:37:25 -0700
Subject: [PATCH] Replace usage of maxSystemThreads with maxSysmsgThreads.

Systrap will never create more stub-worker threads than
maxSysmsgThreads, so there's no need to allocate memory that
will never be used.

Also change initialization of maxSysmsgThreads to happen at
platform creation time, to account for the fact that gVisor
sets GOMAXPROCS at boot.

PiperOrigin-RevId: 633420808
---
 pkg/sentry/platform/systrap/stub_unsafe.go |  2 +-
 pkg/sentry/platform/systrap/subprocess.go  | 17 +++++++++--------
 pkg/sentry/platform/systrap/systrap.go     |  4 ++++
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/pkg/sentry/platform/systrap/stub_unsafe.go b/pkg/sentry/platform/systrap/stub_unsafe.go
index 25a119aa7..462e66923 100644
--- a/pkg/sentry/platform/systrap/stub_unsafe.go
+++ b/pkg/sentry/platform/systrap/stub_unsafe.go
@@ -149,7 +149,7 @@ func stubInit() {
 	// Allocate maxGuestThreads plus ONE because each per-thread stack
 	// has to be aligned to sysmsg.PerThreadMemSize.
 	// Look at sysmsg/sighandler.c:sysmsg_addr() for more details.
-	mapLen, _ = hostarch.PageRoundUp(mapLen + sysmsg.PerThreadMemSize*(maxSystemThreads+1))
+	mapLen, _ = hostarch.PageRoundUp(mapLen + sysmsg.PerThreadMemSize*(uintptr(maxChildThreads+1)))
 
 	// Allocate context queue region
 	stubContextQueueRegion = mapLen
diff --git a/pkg/sentry/platform/systrap/subprocess.go b/pkg/sentry/platform/systrap/subprocess.go
index d978c69f1..d3ef17d76 100644
--- a/pkg/sentry/platform/systrap/subprocess.go
+++ b/pkg/sentry/platform/systrap/subprocess.go
@@ -101,15 +101,16 @@ type requestStub struct {
 	done chan *thread
 }
 
-// maxSysmsgThreads specifies the maximum number of system threads that a
-// subprocess can create in context decoupled mode.
-// TODO(b/268366549): Replace maxSystemThreads below.
-var maxSysmsgThreads = runtime.GOMAXPROCS(0)
+// maxSysmsgThreads is the maximum number of sysmsg threads that a subprocess
+// can create. It is based on GOMAXPROCS and set once, so it must be set after
+// GOMAXPROCS has been adjusted (see loader.go:Args.NumCPU).
+var maxSysmsgThreads = 0
+
+// maxChildThreads is the max number of all child system threads that a
+// subprocess can create, including sysmsg threads.
+var maxChildThreads = 0
 
 const (
-	// maxSystemThreads specifies the maximum number of system threads that a
-	// subprocess may create in order to process the contexts.
-	maxSystemThreads = 4096
 	// maxGuestContexts specifies the maximum number of task contexts that a
 	// subprocess can handle.
 	maxGuestContexts = 4095
@@ -318,7 +319,7 @@ func newSubprocess(create func() (*thread, error), memoryFile *pgalloc.MemoryFil
 	sp := &subprocess{
 		requests:          requests,
 		faultedContexts:   make(map[*platformContext]struct{}),
-		sysmsgStackPool:   pool.Pool{Start: 0, Limit: maxSystemThreads},
+		sysmsgStackPool:   pool.Pool{Start: 0, Limit: uint64(maxChildThreads)},
 		threadContextPool: pool.Pool{Start: 0, Limit: maxGuestContexts},
 		memoryFile:        memoryFile,
 		sysmsgThreads:     make(map[uint32]*sysmsgThread),
diff --git a/pkg/sentry/platform/systrap/systrap.go b/pkg/sentry/platform/systrap/systrap.go
index 7a5310f7a..5527c09b7 100644
--- a/pkg/sentry/platform/systrap/systrap.go
+++ b/pkg/sentry/platform/systrap/systrap.go
@@ -318,6 +318,10 @@ func (*Systrap) MinUserAddress() hostarch.Addr {
 func New() (*Systrap, error) {
 	// CPUID information has been initialized at this point.
 	archState.Init()
+	// GOMAXPROCS has been set at this point.
+	maxSysmsgThreads = runtime.GOMAXPROCS(0)
+	// Account for syscall thread.
+	maxChildThreads = maxSysmsgThreads + 1
 
 	mf, err := createMemoryFile()
 	if err != nil {