runtime/secret: make tests more sturdy

Technically, the garbage collector can take an arbitrary long time until
it finds a given value unreachable. Account for this fact in our zeroing
tests.

Updates #76586.

Change-Id: Ieaf3cec99afb9204a32524ea559c5f1a280a59e6
Reviewed-on: https://go-review.googlesource.com/c/go/+/732980
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Cherry Mui <cherryyz@google.com>
Reviewed-by: Filippo Valsorda <filippo@golang.org>
Auto-Submit: Keith Randall <khr@golang.org>
Reviewed-by: Keith Randall <khr@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

1 files changed, 29 insertions, 9 deletions

diff --git a/src/runtime/secret/secret_test.go b/src/runtime/secret/secret_test.go
index 98d67cf8a4..e2f78c53a0 100644
--- a/src/runtime/secret/secret_test.go
+++ b/src/runtime/secret/secret_test.go
@@ -19,6 +19,7 @@ import (
 	"testing"
 	"time"
 	"unsafe"
+	"weak"
 )
 
 type secretType int64
@@ -63,28 +64,33 @@ func heapSTiny() *secretType {
 // are freed.
 // See runtime/mheap.go:freeSpecial.
 func TestHeap(t *testing.T) {
-	var u uintptr
+	var addr uintptr
+	var p weak.Pointer[S]
 	Do(func() {
-		u = uintptr(unsafe.Pointer(heapS()))
+		sp := heapS()
+		addr = uintptr(unsafe.Pointer(sp))
+		p = weak.Make(sp)
 	})
-
-	runtime.GC()
+	waitCollected(t, p)
 
 	// Check that object got zeroed.
-	checkRangeForSecret(t, u, u+unsafe.Sizeof(S{}))
+	checkRangeForSecret(t, addr, addr+unsafe.Sizeof(S{}))
 	// Also check our stack, just because we can.
 	checkStackForSecret(t)
 }
 
 func TestHeapTiny(t *testing.T) {
-	var u uintptr
+	var addr uintptr
+	var p weak.Pointer[secretType]
 	Do(func() {
-		u = uintptr(unsafe.Pointer(heapSTiny()))
+		sp := heapSTiny()
+		addr = uintptr(unsafe.Pointer(sp))
+		p = weak.Make(sp)
 	})
-	runtime.GC()
+	waitCollected(t, p)
 
 	// Check that object got zeroed.
-	checkRangeForSecret(t, u, u+unsafe.Sizeof(secretType(0)))
+	checkRangeForSecret(t, addr, addr+unsafe.Sizeof(secretType(0)))
 	// Also check our stack, just because we can.
 	checkStackForSecret(t)
 }
@@ -240,6 +246,20 @@ func checkRangeForSecret(t *testing.T, lo, hi uintptr) {
 	}
 }
 
+func waitCollected[P any](t *testing.T, ptr weak.Pointer[P]) {
+	t.Helper()
+	i := 0
+	for ptr.Value() != nil {
+		runtime.GC()
+		i++
+		// 20 seems like a decent number of times to try
+		if i > 20 {
+			t.Errorf("value was never collected")
+		}
+	}
+	t.Logf("number of cycles until collection: %d", i)
+}
+
 func TestRegisters(t *testing.T) {
 	Do(func() {
 		s := makeS()