From 0620c3b6388471f83a1bd36f5a375c03ef68c768 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Tue, 7 Mar 2023 18:25:06 -0800
Subject: [PATCH] cpuid: Initialize `hostFeatureSet` only when it's needed.

This moves the initialization of `cpuid.hostFeatureSet` out of package-level
`init` and instead moves it to an explicit `cpuid.Initialize()` function.

On AMD64, this saves about 512KiB of heap memory that would otherwise always
be live.

PiperOrigin-RevId: 514896323
---
 pkg/cpuid/BUILD                          |  1 +
 pkg/cpuid/cpuid.go                       |  9 +++++++++
 pkg/cpuid/cpuid_parse_test.go            |  6 ++++++
 pkg/cpuid/cpuid_test.go                  |  5 +++++
 pkg/cpuid/native_amd64.go                |  3 ++-
 pkg/cpuid/native_arm64.go                |  3 ++-
 pkg/ring0/lib_amd64.go                   | 11 ++++++-----
 pkg/ring0/lib_arm64.go                   |  4 ++++
 pkg/sentry/fsimpl/testutil/kernel.go     |  1 +
 pkg/sentry/platform/kvm/kvm_amd64.go     |  1 +
 pkg/sentry/platform/kvm/kvm_test.go      | 11 ++++++++++-
 pkg/sentry/platform/kvm/machine_amd64.go |  9 +++++++--
 runsc/boot/BUILD                         |  2 ++
 runsc/boot/loader_test.go                |  8 ++++++++
 runsc/cmd/BUILD                          |  2 ++
 runsc/cmd/boot.go                        |  8 ++++++++
 16 files changed, 74 insertions(+), 10 deletions(-)

diff --git a/pkg/cpuid/BUILD b/pkg/cpuid/BUILD
index eb67c5463..0dc4c26a3 100644
--- a/pkg/cpuid/BUILD
+++ b/pkg/cpuid/BUILD
@@ -21,6 +21,7 @@ go_library(
     visibility = ["//:sandbox"],
     deps = [
         "//pkg/log",
+        "//pkg/sync",
     ],
 )
 
diff --git a/pkg/cpuid/cpuid.go b/pkg/cpuid/cpuid.go
index fae7d277f..413b84930 100644
--- a/pkg/cpuid/cpuid.go
+++ b/pkg/cpuid/cpuid.go
@@ -35,6 +35,7 @@ import (
 	"strings"
 
 	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sync"
 )
 
 // contextID is the package for context.Context.Value keys.
@@ -253,3 +254,11 @@ func initHWCap() {
 		hostFeatureSet.hwCap = c
 	}
 }
+
+var initOnce sync.Once
+
+// Initialize initializes the global data structures used by this package.
+// Must be called prior to using anything else in this package.
+func Initialize() {
+	initOnce.Do(archInitialize)
+}
diff --git a/pkg/cpuid/cpuid_parse_test.go b/pkg/cpuid/cpuid_parse_test.go
index bf1a4e284..d51882948 100644
--- a/pkg/cpuid/cpuid_parse_test.go
+++ b/pkg/cpuid/cpuid_parse_test.go
@@ -16,6 +16,7 @@ package cpuid
 
 import (
 	"io/ioutil"
+	"os"
 	"regexp"
 	"strings"
 	"testing"
@@ -67,3 +68,8 @@ func TestHostFeatureFlags(t *testing.T) {
 		}
 	}
 }
+
+func TestMain(m *testing.M) {
+	Initialize()
+	os.Exit(m.Run())
+}
diff --git a/pkg/cpuid/cpuid_test.go b/pkg/cpuid/cpuid_test.go
index 3ef45d6e4..4281c65fd 100644
--- a/pkg/cpuid/cpuid_test.go
+++ b/pkg/cpuid/cpuid_test.go
@@ -90,3 +90,8 @@ func TestReadingSelfProcAuxv(t *testing.T) {
 		t.Errorf("got %v, expected nil", err)
 	}
 }
+
+func TestMain(m *testing.M) {
+	Initialize()
+	os.Exit(m.Run())
+}
diff --git a/pkg/cpuid/native_amd64.go b/pkg/cpuid/native_amd64.go
index 3bb097ffb..eaf77511d 100644
--- a/pkg/cpuid/native_amd64.go
+++ b/pkg/cpuid/native_amd64.go
@@ -215,7 +215,8 @@ func readMaxCPUFreq() {
 
 }
 
-func init() {
+// archInitialize initializes hostFeatureSet.
+func archInitialize() {
 	hostFeatureSet = FeatureSet{
 		Function: &Native{},
 	}.Fixed()
diff --git a/pkg/cpuid/native_arm64.go b/pkg/cpuid/native_arm64.go
index f4594a8f3..f09edcece 100644
--- a/pkg/cpuid/native_arm64.go
+++ b/pkg/cpuid/native_arm64.go
@@ -150,7 +150,8 @@ func initCPUInfo() {
 	}
 }
 
-func init() {
+// archInitialize initializes hostFeatureSet.
+func archInitialize() {
 	initCPUInfo()
 	initHWCap()
 }
diff --git a/pkg/ring0/lib_amd64.go b/pkg/ring0/lib_amd64.go
index ccd183346..585a5fa2a 100644
--- a/pkg/ring0/lib_amd64.go
+++ b/pkg/ring0/lib_amd64.go
@@ -93,9 +93,9 @@ var (
 
 // Init sets function pointers based on architectural features.
 //
-// This must be called prior to using ring0. By default, it will be called by
-// the init() function. However, it may be called at another time with a
-// different FeatureSet.
+// This must be called prior to using ring0. It may be called with the
+// auto-detected feature set using InitDefault. It may also be called at
+// another time with a  different FeatureSet.
 func Init(fs cpuid.FeatureSet) {
 	// Initialize all sizes.
 	VirtualAddressBits = uintptr(fs.VirtualAddressBits())
@@ -123,7 +123,8 @@ func Init(fs cpuid.FeatureSet) {
 	}
 }
 
-func init() {
-	// See Init, above.
+// InitDefault initializes ring0 with the auto-detected host feature set.
+func InitDefault() {
+	cpuid.Initialize()
 	Init(cpuid.HostFeatureSet())
 }
diff --git a/pkg/ring0/lib_arm64.go b/pkg/ring0/lib_arm64.go
index a72a6926d..3d0a26b7a 100644
--- a/pkg/ring0/lib_arm64.go
+++ b/pkg/ring0/lib_arm64.go
@@ -76,3 +76,7 @@ func FPSIMDEnableTrap()
 //
 // This must be called prior to using ring0.
 func Init() {}
+
+// InitDefault calls Init with default parameters.
+// On ARM, this is not much.
+func InitDefault() {}
diff --git a/pkg/sentry/fsimpl/testutil/kernel.go b/pkg/sentry/fsimpl/testutil/kernel.go
index 420593808..1565fa591 100644
--- a/pkg/sentry/fsimpl/testutil/kernel.go
+++ b/pkg/sentry/fsimpl/testutil/kernel.go
@@ -50,6 +50,7 @@ var (
 
 // Boot initializes a new bare bones kernel for test.
 func Boot() (*kernel.Kernel, error) {
+	cpuid.Initialize()
 	seccheck.Initialize()
 
 	platformCtr, err := platform.Lookup(*platformFlag)
diff --git a/pkg/sentry/platform/kvm/kvm_amd64.go b/pkg/sentry/platform/kvm/kvm_amd64.go
index b7c879163..c61cdf9e0 100644
--- a/pkg/sentry/platform/kvm/kvm_amd64.go
+++ b/pkg/sentry/platform/kvm/kvm_amd64.go
@@ -219,6 +219,7 @@ func (c *cpuidEntries) Set(in cpuid.In, out cpuid.Out) {
 
 // updateGlobalOnce does global initialization. It has to be called only once.
 func updateGlobalOnce(fd int) error {
+	bitsForScaling = getBitsForScaling()
 	if err := updateSystemValues(int(fd)); err != nil {
 		return err
 	}
diff --git a/pkg/sentry/platform/kvm/kvm_test.go b/pkg/sentry/platform/kvm/kvm_test.go
index e82a9e3d3..f0ba03b08 100644
--- a/pkg/sentry/platform/kvm/kvm_test.go
+++ b/pkg/sentry/platform/kvm/kvm_test.go
@@ -16,12 +16,14 @@ package kvm
 
 import (
 	"math/rand"
+	"os"
 	"reflect"
 	"testing"
 	"time"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
+	"gvisor.dev/gvisor/pkg/cpuid"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/ring0"
 	"gvisor.dev/gvisor/pkg/ring0/pagetables"
@@ -32,7 +34,8 @@ import (
 	ktime "gvisor.dev/gvisor/pkg/sentry/time"
 )
 
-var dummyFPState = fpu.NewState()
+// dummyFPState is initialized in TestMain.
+var dummyFPState fpu.State
 
 type testHarness interface {
 	Errorf(format string, args ...any)
@@ -554,3 +557,9 @@ func BenchmarkWorldSwitchToUserRoundtrip(b *testing.B) {
 		b.Logf("ErrContextInterrupt occurred %d times (in %d iterations).", a, a+i)
 	}
 }
+
+func TestMain(m *testing.M) {
+	cpuid.Initialize()
+	dummyFPState = fpu.NewState()
+	os.Exit(m.Run())
+}
diff --git a/pkg/sentry/platform/kvm/machine_amd64.go b/pkg/sentry/platform/kvm/machine_amd64.go
index 5e8e3b511..5209118e2 100644
--- a/pkg/sentry/platform/kvm/machine_amd64.go
+++ b/pkg/sentry/platform/kvm/machine_amd64.go
@@ -158,12 +158,17 @@ func (c *vCPU) initArchState() error {
 }
 
 // bitsForScaling returns the bits available for storing the fraction component
+// of the TSC scaling ratio.
+// It is set using getBitsForScaling when the KVM platform is initialized.
+var bitsForScaling int64
+
+// getBitsForScaling returns the bits available for storing the fraction component
 // of the TSC scaling ratio. This allows us to replicate the (bad) math done by
 // the kernel below in scaledTSC, and ensure we can compute an exact zero
 // offset in setSystemTime.
 //
 // These constants correspond to kvm_tsc_scaling_ratio_frac_bits.
-var bitsForScaling = func() int64 {
+func getBitsForScaling() int64 {
 	fs := cpuid.HostFeatureSet()
 	if fs.Intel() {
 		return 48 // See vmx.c (kvm sources).
@@ -172,7 +177,7 @@ var bitsForScaling = func() int64 {
 	} else {
 		return 63 // Unknown: theoretical maximum.
 	}
-}()
+}
 
 // scaledTSC returns the host TSC scaled by the given frequency.
 //
diff --git a/runsc/boot/BUILD b/runsc/boot/BUILD
index 4dce80c07..1a6f65300 100644
--- a/runsc/boot/BUILD
+++ b/runsc/boot/BUILD
@@ -137,8 +137,10 @@ go_test(
     library = ":boot",
     deps = [
         "//pkg/control/server",
+        "//pkg/cpuid",
         "//pkg/fspath",
         "//pkg/log",
+        "//pkg/sentry/seccheck",
         "//pkg/sentry/vfs",
         "//pkg/sync",
         "//pkg/unet",
diff --git a/runsc/boot/loader_test.go b/runsc/boot/loader_test.go
index a87a30928..a0a8a2fd7 100644
--- a/runsc/boot/loader_test.go
+++ b/runsc/boot/loader_test.go
@@ -26,8 +26,10 @@ import (
 	"github.com/syndtr/gocapability/capability"
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/control/server"
+	"gvisor.dev/gvisor/pkg/cpuid"
 	"gvisor.dev/gvisor/pkg/fspath"
 	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sentry/seccheck"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/unet"
@@ -496,3 +498,9 @@ func TestCreateMountNamespace(t *testing.T) {
 		})
 	}
 }
+
+func TestMain(m *testing.M) {
+	cpuid.Initialize()
+	seccheck.Initialize()
+	os.Exit(m.Run())
+}
diff --git a/runsc/cmd/BUILD b/runsc/cmd/BUILD
index 2d0eb4890..0dde688ff 100644
--- a/runsc/cmd/BUILD
+++ b/runsc/cmd/BUILD
@@ -54,9 +54,11 @@ go_library(
         "//pkg/atomicbitops",
         "//pkg/coretag",
         "//pkg/coverage",
+        "//pkg/cpuid",
         "//pkg/log",
         "//pkg/metric",
         "//pkg/prometheus",
+        "//pkg/ring0",
         "//pkg/sentry/control",
         "//pkg/sentry/fsutil/chdir",
         "//pkg/sentry/kernel",
diff --git a/runsc/cmd/boot.go b/runsc/cmd/boot.go
index 9e69e5ce1..6dfdff5b4 100644
--- a/runsc/cmd/boot.go
+++ b/runsc/cmd/boot.go
@@ -28,8 +28,10 @@ import (
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/coretag"
+	"gvisor.dev/gvisor/pkg/cpuid"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/metric"
+	"gvisor.dev/gvisor/pkg/ring0"
 	"gvisor.dev/gvisor/pkg/sentry/fsutil/chdir"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
 	"gvisor.dev/gvisor/runsc/boot"
@@ -208,6 +210,12 @@ func (b *Boot) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomma
 		}
 	}
 
+	// Initialize CPUID information.
+	cpuid.Initialize()
+
+	// Initialize ring0 library.
+	ring0.InitDefault()
+
 	if len(b.productName) == 0 {
 		// Do this before chroot takes effect, otherwise we can't read /sys.
 		if product, err := ioutil.ReadFile("/sys/devices/virtual/dmi/id/product_name"); err != nil {