runtime, internal/synctest, sync: associate WaitGroups with bubbles

Add support to internal/synctest for managing associations between
arbitrary pointers and synctest bubbles. (Implemented internally to
the runtime package by attaching a special to the pointer.)

Associate WaitGroups with bubbles.
Since WaitGroups don't have a constructor,
perform the association when Add is called.
All Add calls must be made from within the same bubble,
or outside any bubble.

When a bubbled goroutine calls WaitGroup.Wait,
the wait is durably blocking iff the WaitGroup is associated
with the current bubble.

Change-Id: I77e2701e734ac2fa2b32b28d5b0c853b7b2825c9
Reviewed-on: https://go-review.googlesource.com/c/go/+/676656
Reviewed-by: Michael Knyszek <mknyszek@google.com>
Reviewed-by: Michael Pratt <mpratt@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Damien Neil <dneil@google.com>

2 files changed, 132 insertions, 2 deletions

diff --git a/src/internal/synctest/synctest.go b/src/internal/synctest/synctest.go
index 19190d30f1..4d7fa3730c 100644
--- a/src/internal/synctest/synctest.go
+++ b/src/internal/synctest/synctest.go
@@ -8,7 +8,7 @@
 package synctest
 
 import (
-	_ "unsafe" // for go:linkname
+	"unsafe"
 )
 
 //go:linkname Run
@@ -17,6 +17,36 @@ func Run(f func())
 //go:linkname Wait
 func Wait()
 
+// IsInBubble reports whether the current goroutine is in a bubble.
+//
+//go:linkname IsInBubble
+func IsInBubble() bool
+
+// Associate associates p with the current bubble.
+// It returns false if p has an existing association with a different bubble.
+func Associate[T any](p *T) (ok bool) {
+	return associate(unsafe.Pointer(p))
+}
+
+//go:linkname associate
+func associate(p unsafe.Pointer) bool
+
+// Disassociate disassociates p from any bubble.
+func Disassociate[T any](p *T) {
+	disassociate(unsafe.Pointer(p))
+}
+
+//go:linkname disassociate
+func disassociate(b unsafe.Pointer)
+
+// IsAssociated reports whether p is associated with the current bubble.
+func IsAssociated[T any](p *T) bool {
+	return isAssociated(unsafe.Pointer(p))
+}
+
+//go:linkname isAssociated
+func isAssociated(p unsafe.Pointer) bool
+
 //go:linkname acquire
 func acquire() any
 
diff --git a/src/internal/synctest/synctest_test.go b/src/internal/synctest/synctest_test.go
index 7f71df1710..8b2ade5630 100644
--- a/src/internal/synctest/synctest_test.go
+++ b/src/internal/synctest/synctest_test.go
@@ -7,11 +7,14 @@ package synctest_test
 import (
 	"fmt"
 	"internal/synctest"
+	"internal/testenv"
 	"iter"
+	"os"
 	"reflect"
 	"runtime"
 	"slices"
 	"strconv"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -523,7 +526,7 @@ func TestReflectFuncOf(t *testing.T) {
 	})
 }
 
-func TestWaitGroup(t *testing.T) {
+func TestWaitGroupInBubble(t *testing.T) {
 	synctest.Run(func() {
 		var wg sync.WaitGroup
 		wg.Add(1)
@@ -540,6 +543,83 @@ func TestWaitGroup(t *testing.T) {
 	})
 }
 
+func TestWaitGroupOutOfBubble(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	donec := make(chan struct{})
+	go synctest.Run(func() {
+		// Since wg.Add was called outside the bubble, Wait is not durably blocking
+		// and this waits until wg.Done is called below.
+		wg.Wait()
+		close(donec)
+	})
+	select {
+	case <-donec:
+		t.Fatalf("synctest.Run finished before WaitGroup.Done called")
+	case <-time.After(1 * time.Millisecond):
+	}
+	wg.Done()
+	<-donec
+}
+
+func TestWaitGroupMovedIntoBubble(t *testing.T) {
+	wantFatal(t, "fatal error: sync: WaitGroup.Add called from inside and outside synctest bubble", func() {
+		var wg sync.WaitGroup
+		wg.Add(1)
+		synctest.Run(func() {
+			wg.Add(1)
+		})
+	})
+}
+
+func TestWaitGroupMovedOutOfBubble(t *testing.T) {
+	wantFatal(t, "fatal error: sync: WaitGroup.Add called from inside and outside synctest bubble", func() {
+		var wg sync.WaitGroup
+		synctest.Run(func() {
+			wg.Add(1)
+		})
+		wg.Add(1)
+	})
+}
+
+func TestWaitGroupMovedBetweenBubblesWithNonZeroCount(t *testing.T) {
+	wantFatal(t, "fatal error: sync: WaitGroup.Add called from multiple synctest bubbles", func() {
+		var wg sync.WaitGroup
+		synctest.Run(func() {
+			wg.Add(1)
+		})
+		synctest.Run(func() {
+			wg.Add(1)
+		})
+	})
+}
+
+func TestWaitGroupMovedBetweenBubblesWithZeroCount(t *testing.T) {
+	var wg sync.WaitGroup
+	synctest.Run(func() {
+		wg.Add(1)
+		wg.Done()
+	})
+	synctest.Run(func() {
+		// Reusing the WaitGroup is safe, because its count is zero.
+		wg.Add(1)
+		wg.Done()
+	})
+}
+
+func TestWaitGroupMovedBetweenBubblesAfterWait(t *testing.T) {
+	var wg sync.WaitGroup
+	synctest.Run(func() {
+		wg.Go(func() {})
+		wg.Wait()
+	})
+	synctest.Run(func() {
+		// Reusing the WaitGroup is safe, because its count is zero.
+		wg.Go(func() {})
+		wg.Wait()
+	})
+}
+
 func TestHappensBefore(t *testing.T) {
 	// Use two parallel goroutines accessing different vars to ensure that
 	// we correctly account for multiple goroutines in the bubble.
@@ -647,3 +727,23 @@ func wantPanic(t *testing.T, want string) {
 		t.Errorf("got no panic, want one")
 	}
 }
+
+func wantFatal(t *testing.T, want string, f func()) {
+	t.Helper()
+
+	if os.Getenv("GO_WANT_HELPER_PROCESS") == "1" {
+		f()
+		return
+	}
+
+	cmd := testenv.Command(t, testenv.Executable(t), "-test.run=^"+t.Name()+"$")
+	cmd = testenv.CleanCmdEnv(cmd)
+	cmd.Env = append(cmd.Env, "GO_WANT_HELPER_PROCESS=1")
+	out, err := cmd.CombinedOutput()
+	if err == nil {
+		t.Errorf("expected test function to panic, but test returned successfully")
+	}
+	if !strings.Contains(string(out), want) {
+		t.Errorf("wanted test output contaiing %q; got %q", want, string(out))
+	}
+}