Allow goid.Get() to be used outside of race builds.

PiperOrigin-RevId: 342179912

6 files changed, 105 insertions, 164 deletions

diff --git a/pkg/goid/BUILD b/pkg/goid/BUILD
index 7a82631c5..d855b702c 100644
--- a/pkg/goid/BUILD
+++ b/pkg/goid/BUILD
@@ -8,8 +8,6 @@ go_library(
         "goid.go",
         "goid_amd64.s",
         "goid_arm64.s",
-        "goid_race.go",
-        "goid_unsafe.go",
     ],
     visibility = ["//visibility:public"],
 )
@@ -18,7 +16,6 @@ go_test(
     name = "goid_test",
     size = "small",
     srcs = [
-        "empty_test.go",
         "goid_test.go",
     ],
     library = ":goid",
diff --git a/pkg/goid/empty_test.go b/pkg/goid/empty_test.go
deleted file mode 100644
index c0a4b17ab..000000000
--- a/pkg/goid/empty_test.go
+++ /dev/null
@@ -1,22 +0,0 @@
-// Copyright 2020 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.
-
-// +build !race
-
-package goid
-
-import "testing"
-
-// TestNothing exists to make the build system happy.
-func TestNothing(t *testing.T) {}
diff --git a/pkg/goid/goid.go b/pkg/goid/goid.go
index 39df30031..17c384cb0 100644
--- a/pkg/goid/goid.go
+++ b/pkg/goid/goid.go
@@ -12,13 +12,61 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// +build !race
+// +build go1.12
+// +build !go1.17
 
-// Package goid provides access to the ID of the current goroutine in
-// race/gotsan builds.
+// Check type signatures when updating Go version.
+
+// Package goid provides the Get function.
 package goid
 
 // Get returns the ID of the current goroutine.
 func Get() int64 {
-	panic("unimplemented for non-race builds")
+	return getg().goid
+}
+
+// Structs from Go runtime. These may change in the future and require
+// updating. These structs are currently the same on both AMD64 and ARM64,
+// but may diverge in the future.
+
+type stack struct {
+	lo uintptr
+	hi uintptr
+}
+
+type gobuf struct {
+	sp   uintptr
+	pc   uintptr
+	g    uintptr
+	ctxt uintptr
+	ret  uint64
+	lr   uintptr
+	bp   uintptr
 }
+
+type g struct {
+	stack       stack
+	stackguard0 uintptr
+	stackguard1 uintptr
+
+	_panic       uintptr
+	_defer       uintptr
+	m            uintptr
+	sched        gobuf
+	syscallsp    uintptr
+	syscallpc    uintptr
+	stktopsp     uintptr
+	param        uintptr
+	atomicstatus uint32
+	stackLock    uint32
+	goid         int64
+
+	// More fields...
+	//
+	// We only use goid and the fields before it are only listed to
+	// calculate the correct offset.
+}
+
+// Defined in assembly. This can't use go:linkname since runtime.getg() isn't a
+// real function, it's a compiler intrinsic.
+func getg() *g
diff --git a/pkg/goid/goid_race.go b/pkg/goid/goid_race.go
deleted file mode 100644
index 1766beaee..000000000
--- a/pkg/goid/goid_race.go
+++ /dev/null
@@ -1,25 +0,0 @@
-// Copyright 2020 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.
-
-// Only available in race/gotsan builds.
-// +build race
-
-// Package goid provides access to the ID of the current goroutine in
-// race/gotsan builds.
-package goid
-
-// Get returns the ID of the current goroutine.
-func Get() int64 {
-	return goid()
-}
diff --git a/pkg/goid/goid_test.go b/pkg/goid/goid_test.go
index 31970ce79..54be11d63 100644
--- a/pkg/goid/goid_test.go
+++ b/pkg/goid/goid_test.go
@@ -12,63 +12,70 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// +build race
-
 package goid
 
 import (
 	"runtime"
 	"sync"
 	"testing"
+	"time"
 )
 
-func TestInitialGoID(t *testing.T) {
-	const max = 10000
-	if id := goid(); id < 0 || id > max {
-		t.Errorf("got goid = %d, want 0 < goid <= %d", id, max)
-	}
-}
+func TestUniquenessAndConsistency(t *testing.T) {
+	const (
+		numGoroutines = 5000
 
-// TestGoIDSquence verifies that goid returns values which could plausibly be
-// goroutine IDs. If this test breaks or becomes flaky, the structs in
-// goid_unsafe.go may need to be updated.
-func TestGoIDSquence(t *testing.T) {
-	// Goroutine IDs are cached by each P.
-	runtime.GOMAXPROCS(1)
+		// maxID is not an intrinsic property of goroutine IDs; it is only a
+		// property of how the Go runtime currently assigns them. Future
+		// changes to the Go runtime may require that maxID be raised, or that
+		// assertions regarding it be removed entirely.
+		maxID = numGoroutines + 1000
+	)
 
-	// Fill any holes in lower range.
-	for i := 0; i < 50; i++ {
-		var wg sync.WaitGroup
-		wg.Add(1)
+	var (
+		goidsMu   sync.Mutex
+		goids     = make(map[int64]struct{})
+		checkedWG sync.WaitGroup
+		exitCh    = make(chan struct{})
+	)
+	for i := 0; i < numGoroutines; i++ {
+		checkedWG.Add(1)
 		go func() {
-			wg.Done()
-
-			// Leak the goroutine to prevent the ID from being
-			// reused.
-			select {}
-		}()
-		wg.Wait()
-	}
-
-	id := goid()
-	for i := 0; i < 100; i++ {
-		var (
-			newID int64
-			wg    sync.WaitGroup
-		)
-		wg.Add(1)
-		go func() {
-			newID = goid()
-			wg.Done()
-
-			// Leak the goroutine to prevent the ID from being
-			// reused.
-			select {}
+			id := Get()
+			if id > maxID {
+				t.Errorf("observed unexpectedly large goroutine ID %d", id)
+			}
+			goidsMu.Lock()
+			if _, dup := goids[id]; dup {
+				t.Errorf("observed duplicate goroutine ID %d", id)
+			}
+			goids[id] = struct{}{}
+			goidsMu.Unlock()
+			checkedWG.Done()
+			for {
+				if curID := Get(); curID != id {
+					t.Errorf("goroutine ID changed from %d to %d", id, curID)
+					// Don't spam logs by repeating the check; wait quietly for
+					// the test to finish.
+					<-exitCh
+					return
+				}
+				// Check if the test is over.
+				select {
+				case <-exitCh:
+					return
+				default:
+				}
+				// Yield to other goroutines, and possibly migrate to another P.
+				runtime.Gosched()
+			}
 		}()
-		wg.Wait()
-		if max := id + 100; newID <= id || newID > max {
-			t.Errorf("unexpected goroutine ID pattern, got goid = %d, want %d < goid <= %d (previous = %d)", newID, id, max, id)
-		}
-		id = newID
 	}
+	// Wait for all goroutines to perform uniqueness checks.
+	checkedWG.Wait()
+	// Wait for an additional second to allow goroutines to spin checking for
+	// ID consistency.
+	time.Sleep(time.Second)
+	// Request that all goroutines exit.
+	close(exitCh)
 }
diff --git a/pkg/goid/goid_unsafe.go b/pkg/goid/goid_unsafe.go
deleted file mode 100644
index ded8004dd..000000000
--- a/pkg/goid/goid_unsafe.go
+++ /dev/null
@@ -1,64 +0,0 @@
-// Copyright 2020 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 goid
-
-// Structs from Go runtime. These may change in the future and require
-// updating. These structs are currently the same on both AMD64 and ARM64,
-// but may diverge in the future.
-
-type stack struct {
-	lo uintptr
-	hi uintptr
-}
-
-type gobuf struct {
-	sp   uintptr
-	pc   uintptr
-	g    uintptr
-	ctxt uintptr
-	ret  uint64
-	lr   uintptr
-	bp   uintptr
-}
-
-type g struct {
-	stack       stack
-	stackguard0 uintptr
-	stackguard1 uintptr
-
-	_panic       uintptr
-	_defer       uintptr
-	m            uintptr
-	sched        gobuf
-	syscallsp    uintptr
-	syscallpc    uintptr
-	stktopsp     uintptr
-	param        uintptr
-	atomicstatus uint32
-	stackLock    uint32
-	goid         int64
-
-	// More fields...
-	//
-	// We only use goid and the fields before it are only listed to
-	// calculate the correct offset.
-}
-
-func getg() *g
-
-// goid returns the ID of the current goroutine.
-func goid() int64 {
-	return getg().goid
-}