From ab32fa2481d84e3f390f2cd81e25c537756b7aa1 Mon Sep 17 00:00:00 2001
From: Adin Scannell <ascannell@google.com>
Date: Tue, 5 Jan 2021 18:13:18 -0800
Subject: [PATCH] Make type sanity checking happen only in race builds.

This adds significant costs to startup, since it is done for
every type in the system. Since the state package already saves
sanity checks for race builds, use this for type registration.

PiperOrigin-RevId: 350259336
---
 pkg/state/tests/register_test.go | 11 +++++
 pkg/state/tests/struct_test.go   | 11 -----
 pkg/state/types.go               | 75 ++++++++++++++++----------------
 3 files changed, 49 insertions(+), 48 deletions(-)

diff --git a/pkg/state/tests/register_test.go b/pkg/state/tests/register_test.go
index c829753cc..75bdbfc6e 100644
--- a/pkg/state/tests/register_test.go
+++ b/pkg/state/tests/register_test.go
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// +build race
+
 package tests
 
 import (
@@ -165,3 +167,12 @@ func TestRegisterBad(t *testing.T) {
 
 	}
 }
+
+func TestRegisterTypeOnlyStruct(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("Register did not panic")
+		}
+	}()
+	state.Register((*typeOnlyEmptyStruct)(nil))
+}
diff --git a/pkg/state/tests/struct_test.go b/pkg/state/tests/struct_test.go
index c91c2c032..9826f1ee9 100644
--- a/pkg/state/tests/struct_test.go
+++ b/pkg/state/tests/struct_test.go
@@ -17,8 +17,6 @@ package tests
 import (
 	"math/rand"
 	"testing"
-
-	"gvisor.dev/gvisor/pkg/state"
 )
 
 func TestEmptyStruct(t *testing.T) {
@@ -58,15 +56,6 @@ func TestEmptyStruct(t *testing.T) {
 	})
 }
 
-func TestRegisterTypeOnlyStruct(t *testing.T) {
-	defer func() {
-		if r := recover(); r == nil {
-			t.Errorf("Register did not panic")
-		}
-	}()
-	state.Register((*typeOnlyEmptyStruct)(nil))
-}
-
 func TestEmbeddedPointers(t *testing.T) {
 	// Give each int64 a random value to prevent Go from using
 	// runtime.staticuint64s, which confounds tests for struct duplication.
diff --git a/pkg/state/types.go b/pkg/state/types.go
index 84aed8732..420675880 100644
--- a/pkg/state/types.go
+++ b/pkg/state/types.go
@@ -329,47 +329,48 @@ var reverseTypeDatabase = map[reflect.Type]string{}
 // This must be called on init and only done once.
 func Register(t Type) {
 	name := t.StateTypeName()
-	fields := t.StateFields()
-	assertValidType(name, fields)
-	// Register must always be called on pointers.
 	typ := reflect.TypeOf(t)
-	if typ.Kind() != reflect.Ptr {
-		Failf("Register must be called on pointers")
+	if raceEnabled {
+		assertValidType(name, t.StateFields())
+		// Register must always be called on pointers.
+		if typ.Kind() != reflect.Ptr {
+			Failf("Register must be called on pointers")
+		}
 	}
 	typ = typ.Elem()
-	if typ.Kind() == reflect.Struct {
-		// All registered structs must implement SaverLoader. We allow
-		// the registration is non-struct types with just the Type
-		// interface, but we need to call StateSave/StateLoad methods
-		// on aggregate types.
-		if _, ok := t.(SaverLoader); !ok {
-			Failf("struct %T does not implement SaverLoader", t)
-		}
-	} else {
-		// Non-structs must not have any fields. We don't support
-		// calling StateSave/StateLoad methods on any non-struct types.
-		// If custom behavior is required, these types should be
-		// wrapped in a structure of some kind.
-		if len(fields) != 0 {
-			Failf("non-struct %T has non-zero fields %v", t, fields)
-		}
-		// We don't allow non-structs to implement StateSave/StateLoad
-		// methods, because they won't be called and it's confusing.
-		if _, ok := t.(SaverLoader); ok {
-			Failf("non-struct %T implements SaverLoader", t)
-		}
-	}
-	if _, ok := primitiveTypeDatabase[name]; ok {
-		Failf("conflicting primitiveTypeDatabase entry for %T: used by primitive", t)
-	}
-	if _, ok := globalTypeDatabase[name]; ok {
-		Failf("conflicting globalTypeDatabase entries for %T: name conflict", t)
-	}
-	if name == interfaceType {
-		Failf("conflicting name for %T: matches interfaceType", t)
-	}
-	globalTypeDatabase[name] = typ
 	if raceEnabled {
+		if typ.Kind() == reflect.Struct {
+			// All registered structs must implement SaverLoader. We allow
+			// the registration is non-struct types with just the Type
+			// interface, but we need to call StateSave/StateLoad methods
+			// on aggregate types.
+			if _, ok := t.(SaverLoader); !ok {
+				Failf("struct %T does not implement SaverLoader", t)
+			}
+		} else {
+			// Non-structs must not have any fields. We don't support
+			// calling StateSave/StateLoad methods on any non-struct types.
+			// If custom behavior is required, these types should be
+			// wrapped in a structure of some kind.
+			if fields := t.StateFields(); len(fields) != 0 {
+				Failf("non-struct %T has non-zero fields %v", t, fields)
+			}
+			// We don't allow non-structs to implement StateSave/StateLoad
+			// methods, because they won't be called and it's confusing.
+			if _, ok := t.(SaverLoader); ok {
+				Failf("non-struct %T implements SaverLoader", t)
+			}
+		}
+		if _, ok := primitiveTypeDatabase[name]; ok {
+			Failf("conflicting primitiveTypeDatabase entry for %T: used by primitive", t)
+		}
+		if _, ok := globalTypeDatabase[name]; ok {
+			Failf("conflicting globalTypeDatabase entries for %T: name conflict", t)
+		}
+		if name == interfaceType {
+			Failf("conflicting name for %T: matches interfaceType", t)
+		}
 		reverseTypeDatabase[typ] = name
 	}
+	globalTypeDatabase[name] = typ
 }