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>

4 files changed, 127 insertions, 20 deletions

diff --git a/src/runtime/mheap.go b/src/runtime/mheap.go
index 3612d71e66..9361089b80 100644
--- a/src/runtime/mheap.go
+++ b/src/runtime/mheap.go
@@ -223,6 +223,7 @@ type mheap struct {
 	specialReachableAlloc      fixalloc // allocator for specialReachable
 	specialPinCounterAlloc     fixalloc // allocator for specialPinCounter
 	specialWeakHandleAlloc     fixalloc // allocator for specialWeakHandle
+	specialBubbleAlloc         fixalloc // allocator for specialBubble
 	speciallock                mutex    // lock for special record allocators.
 	arenaHintAlloc             fixalloc // allocator for arenaHints
 
@@ -799,6 +800,7 @@ func (h *mheap) init() {
 	h.specialReachableAlloc.init(unsafe.Sizeof(specialReachable{}), nil, nil, &memstats.other_sys)
 	h.specialPinCounterAlloc.init(unsafe.Sizeof(specialPinCounter{}), nil, nil, &memstats.other_sys)
 	h.specialWeakHandleAlloc.init(unsafe.Sizeof(specialWeakHandle{}), nil, nil, &memstats.gcMiscSys)
+	h.specialBubbleAlloc.init(unsafe.Sizeof(specialBubble{}), nil, nil, &memstats.other_sys)
 	h.arenaHintAlloc.init(unsafe.Sizeof(arenaHint{}), nil, nil, &memstats.other_sys)
 
 	// Don't zero mspan allocations. Background sweeping can
@@ -2003,6 +2005,8 @@ const (
 	// _KindSpecialCheckFinalizer adds additional context to a finalizer or cleanup.
 	// Used only if debug.checkfinalizers != 0.
 	_KindSpecialCheckFinalizer = 8
+	// _KindSpecialBubble is used to associate objects with synctest bubbles.
+	_KindSpecialBubble = 9
 )
 
 type special struct {
@@ -2839,6 +2843,11 @@ func freeSpecial(s *special, p unsafe.Pointer, size uintptr) {
 		lock(&mheap_.speciallock)
 		mheap_.specialTinyBlockAlloc.free(unsafe.Pointer(st))
 		unlock(&mheap_.speciallock)
+	case _KindSpecialBubble:
+		st := (*specialBubble)(unsafe.Pointer(s))
+		lock(&mheap_.speciallock)
+		mheap_.specialBubbleAlloc.free(unsafe.Pointer(st))
+		unlock(&mheap_.speciallock)
 	default:
 		throw("bad special kind")
 		panic("not reached")
diff --git a/src/runtime/runtime2.go b/src/runtime/runtime2.go
index 94ab87f6db..cd40586bc2 100644
--- a/src/runtime/runtime2.go
+++ b/src/runtime/runtime2.go
@@ -1096,6 +1096,7 @@ const (
 	waitReasonSynctestChanReceive                     // "chan receive (synctest)"
 	waitReasonSynctestChanSend                        // "chan send (synctest)"
 	waitReasonSynctestSelect                          // "select (synctest)"
+	waitReasonSynctestWaitGroupWait                   // "sync.WaitGroup.Wait (synctest)"
 	waitReasonCleanupWait                             // "cleanup wait"
 )
 
@@ -1145,6 +1146,7 @@ var waitReasonStrings = [...]string{
 	waitReasonSynctestChanReceive:   "chan receive (synctest)",
 	waitReasonSynctestChanSend:      "chan send (synctest)",
 	waitReasonSynctestSelect:        "select (synctest)",
+	waitReasonSynctestWaitGroupWait: "sync.WaitGroup.Wait (synctest)",
 	waitReasonCleanupWait:           "cleanup wait",
 }
 
@@ -1190,18 +1192,18 @@ func (w waitReason) isIdleInSynctest() bool {
 
 // isIdleInSynctest indicates that a goroutine is considered idle by synctest.Wait.
 var isIdleInSynctest = [len(waitReasonStrings)]bool{
-	waitReasonChanReceiveNilChan:  true,
-	waitReasonChanSendNilChan:     true,
-	waitReasonSelectNoCases:       true,
-	waitReasonSleep:               true,
-	waitReasonSyncCondWait:        true,
-	waitReasonSyncWaitGroupWait:   true,
-	waitReasonCoroutine:           true,
-	waitReasonSynctestRun:         true,
-	waitReasonSynctestWait:        true,
-	waitReasonSynctestChanReceive: true,
-	waitReasonSynctestChanSend:    true,
-	waitReasonSynctestSelect:      true,
+	waitReasonChanReceiveNilChan:    true,
+	waitReasonChanSendNilChan:       true,
+	waitReasonSelectNoCases:         true,
+	waitReasonSleep:                 true,
+	waitReasonSyncCondWait:          true,
+	waitReasonSynctestWaitGroupWait: true,
+	waitReasonCoroutine:             true,
+	waitReasonSynctestRun:           true,
+	waitReasonSynctestWait:          true,
+	waitReasonSynctestChanReceive:   true,
+	waitReasonSynctestChanSend:      true,
+	waitReasonSynctestSelect:        true,
 }
 
 var (
diff --git a/src/runtime/sema.go b/src/runtime/sema.go
index 4890df3464..7d6fc6d57d 100644
--- a/src/runtime/sema.go
+++ b/src/runtime/sema.go
@@ -106,8 +106,12 @@ func sync_runtime_SemacquireRWMutex(addr *uint32, lifo bool, skipframes int) {
 }
 
 //go:linkname sync_runtime_SemacquireWaitGroup sync.runtime_SemacquireWaitGroup
-func sync_runtime_SemacquireWaitGroup(addr *uint32) {
-	semacquire1(addr, false, semaBlockProfile, 0, waitReasonSyncWaitGroupWait)
+func sync_runtime_SemacquireWaitGroup(addr *uint32, synctestDurable bool) {
+	reason := waitReasonSyncWaitGroupWait
+	if synctestDurable {
+		reason = waitReasonSynctestWaitGroupWait
+	}
+	semacquire1(addr, false, semaBlockProfile, 0, reason)
 }
 
 //go:linkname poll_runtime_Semrelease internal/poll.runtime_Semrelease
diff --git a/src/runtime/synctest.go b/src/runtime/synctest.go
index ff1979a8d8..f676afa20d 100644
--- a/src/runtime/synctest.go
+++ b/src/runtime/synctest.go
@@ -5,6 +5,7 @@
 package runtime
 
 import (
+	"internal/runtime/atomic"
 	"internal/runtime/sys"
 	"unsafe"
 )
@@ -13,12 +14,13 @@ import (
 type synctestBubble struct {
 	mu      mutex
 	timers  timers
-	now     int64 // current fake time
-	root    *g    // caller of synctest.Run
-	waiter  *g    // caller of synctest.Wait
-	main    *g    // goroutine started by synctest.Run
-	waiting bool  // true if a goroutine is calling synctest.Wait
-	done    bool  // true if main has exited
+	id      uint64 // unique id
+	now     int64  // current fake time
+	root    *g     // caller of synctest.Run
+	waiter  *g     // caller of synctest.Wait
+	main    *g     // goroutine started by synctest.Run
+	waiting bool   // true if a goroutine is calling synctest.Wait
+	done    bool   // true if main has exited
 
 	// The bubble is active (not blocked) so long as running > 0 || active > 0.
 	//
@@ -163,6 +165,8 @@ func (bubble *synctestBubble) raceaddr() unsafe.Pointer {
 	return unsafe.Pointer(bubble)
 }
 
+var bubbleGen atomic.Uint64 // bubble ID counter
+
 //go:linkname synctestRun internal/synctest.Run
 func synctestRun(f func()) {
 	if debug.asynctimerchan.Load() != 0 {
@@ -174,6 +178,7 @@ func synctestRun(f func()) {
 		panic("synctest.Run called from within a synctest bubble")
 	}
 	bubble := &synctestBubble{
+		id:      bubbleGen.Add(1),
 		total:   1,
 		running: 1,
 		root:    gp,
@@ -313,6 +318,11 @@ func synctestwait_c(gp *g, _ unsafe.Pointer) bool {
 	return true
 }
 
+//go:linkname synctest_isInBubble internal/synctest.IsInBubble
+func synctest_isInBubble() bool {
+	return getg().bubble != nil
+}
+
 //go:linkname synctest_acquire internal/synctest.acquire
 func synctest_acquire() any {
 	if bubble := getg().bubble; bubble != nil {
@@ -339,3 +349,85 @@ func synctest_inBubble(bubble any, f func()) {
 	}()
 	f()
 }
+
+// specialBubble is a special used to associate objects with bubbles.
+type specialBubble struct {
+	_        sys.NotInHeap
+	special  special
+	bubbleid uint64
+}
+
+// getOrSetBubbleSpecial checks the special record for p's bubble membership.
+//
+// If add is true and p is not associated with any bubble,
+// it adds a special record for p associating it with bubbleid.
+//
+// It returns ok==true if p is associated with bubbleid
+// (including if a new association was added),
+// and ok==false if not.
+func getOrSetBubbleSpecial(p unsafe.Pointer, bubbleid uint64, add bool) (ok bool) {
+	span := spanOfHeap(uintptr(p))
+	if span == nil {
+		throw("getOrSetBubbleSpecial on invalid pointer")
+	}
+
+	// Ensure that the span is swept.
+	// Sweeping accesses the specials list w/o locks, so we have
+	// to synchronize with it. And it's just much safer.
+	mp := acquirem()
+	span.ensureSwept()
+
+	offset := uintptr(p) - span.base()
+
+	lock(&span.speciallock)
+
+	// Find splice point, check for existing record.
+	iter, exists := span.specialFindSplicePoint(offset, _KindSpecialBubble)
+	if exists {
+		// p is already associated with a bubble.
+		// Return true iff it's the same bubble.
+		s := (*specialBubble)((unsafe.Pointer)(*iter))
+		ok = s.bubbleid == bubbleid
+	} else if add {
+		// p is not associated with a bubble,
+		// and we've been asked to add an association.
+		s := (*specialBubble)(mheap_.specialBubbleAlloc.alloc())
+		s.bubbleid = bubbleid
+		s.special.kind = _KindSpecialBubble
+		s.special.offset = offset
+		s.special.next = *iter
+		*iter = (*special)(unsafe.Pointer(s))
+		spanHasSpecials(span)
+		ok = true
+	} else {
+		// p is not associated with a bubble.
+		ok = false
+	}
+
+	unlock(&span.speciallock)
+	releasem(mp)
+
+	return ok
+}
+
+// synctest_associate associates p with the current bubble.
+// It returns false if p is already associated with a different bubble.
+//
+//go:linkname synctest_associate internal/synctest.associate
+func synctest_associate(p unsafe.Pointer) (ok bool) {
+	return getOrSetBubbleSpecial(p, getg().bubble.id, true)
+}
+
+// synctest_disassociate disassociates p from its bubble.
+//
+//go:linkname synctest_disassociate internal/synctest.disassociate
+func synctest_disassociate(p unsafe.Pointer) {
+	removespecial(p, _KindSpecialBubble)
+}
+
+// synctest_isAssociated reports whether p is associated with the current bubble.
+//
+//go:linkname synctest_isAssociated internal/synctest.isAssociated
+func synctest_isAssociated(p unsafe.Pointer) bool {
+	return getOrSetBubbleSpecial(p, getg().bubble.id, false)
+}