runtime/internal/atomic: add write barrier-enabled pointer atomics

UnsafePointer.Store, UnsafePointer.CompareAndSwap were missing,
although .StoreNoWB and .CompareAndSwapNoWB existed.
Same for Pointer[T}.

Do the linkname tricks necessary to add those methods.

Change-Id: I925ee27673288accb15ebe93898f9eb01ab46a98
Reviewed-on: https://go-review.googlesource.com/c/go/+/443379
Auto-Submit: Russ Cox <rsc@golang.org>
Reviewed-by: Austin Clements <austin@google.com>
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>

2 files changed, 65 insertions, 16 deletions

diff --git a/src/runtime/atomic_pointer.go b/src/runtime/atomic_pointer.go
index b8f0c22c63..25e0e651b4 100644
--- a/src/runtime/atomic_pointer.go
+++ b/src/runtime/atomic_pointer.go
@@ -35,6 +35,27 @@ func atomicstorep(ptr unsafe.Pointer, new unsafe.Pointer) {
 	atomic.StorepNoWB(noescape(ptr), new)
 }
 
+// atomic_storePointer is the implementation of runtime/internal/UnsafePointer.Store
+// (like StoreNoWB but with the write barrier).
+//
+//go:nosplit
+//go:linkname atomic_storePointer runtime/internal/atomic.storePointer
+func atomic_storePointer(ptr *unsafe.Pointer, new unsafe.Pointer) {
+	atomicstorep(unsafe.Pointer(ptr), new)
+}
+
+// atomic_casPointer is the implementation of runtime/internal/UnsafePointer.CompareAndSwap
+// (like CompareAndSwapNoWB but with the write barrier).
+//
+//go:nosplit
+//go:linkname atomic_casPointer runtime/internal/atomic.casPointer
+func atomic_casPointer(ptr *unsafe.Pointer, old, new unsafe.Pointer) bool {
+	if writeBarrier.enabled {
+		atomicwb(ptr, new)
+	}
+	return atomic.Casp1(ptr, old, new)
+}
+
 // Like above, but implement in terms of sync/atomic's uintptr operations.
 // We cannot just call the runtime routines, because the race detector expects
 // to be able to intercept the sync/atomic forms but not the runtime forms.
diff --git a/src/runtime/internal/atomic/types.go b/src/runtime/internal/atomic/types.go
index 35d8935c70..0d75226b19 100644
--- a/src/runtime/internal/atomic/types.go
+++ b/src/runtime/internal/atomic/types.go
@@ -30,8 +30,7 @@ func (i *Int32) Store(value int32) {
 
 // CompareAndSwap atomically compares i's value with old,
 // and if they're equal, swaps i's value with new.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 //go:nosplit
 func (i *Int32) CompareAndSwap(old, new int32) bool {
@@ -84,8 +83,7 @@ func (i *Int64) Store(value int64) {
 
 // CompareAndSwap atomically compares i's value with old,
 // and if they're equal, swaps i's value with new.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 //go:nosplit
 func (i *Int64) CompareAndSwap(old, new int64) bool {
@@ -231,8 +229,7 @@ func (u *Uint32) StoreRelease(value uint32) {
 
 // CompareAndSwap atomically compares u's value with old,
 // and if they're equal, swaps u's value with new.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 //go:nosplit
 func (u *Uint32) CompareAndSwap(old, new uint32) bool {
@@ -244,8 +241,7 @@ func (u *Uint32) CompareAndSwap(old, new uint32) bool {
 // may observe operations that occur after this operation to
 // precede it, but no operation that precedes it
 // on this thread can be observed to occur after it.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 // WARNING: Use sparingly and with great care.
 //
@@ -322,8 +318,7 @@ func (u *Uint64) Store(value uint64) {
 
 // CompareAndSwap atomically compares u's value with old,
 // and if they're equal, swaps u's value with new.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 //go:nosplit
 func (u *Uint64) CompareAndSwap(old, new uint64) bool {
@@ -399,8 +394,7 @@ func (u *Uintptr) StoreRelease(value uintptr) {
 
 // CompareAndSwap atomically compares u's value with old,
 // and if they're equal, swaps u's value with new.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 //go:nosplit
 func (u *Uintptr) CompareAndSwap(old, new uintptr) bool {
@@ -478,28 +472,47 @@ func (u *UnsafePointer) Load() unsafe.Pointer {
 // perform a write barrier on value, and so this operation may
 // hide pointers from the GC. Use with care and sparingly.
 // It is safe to use with values not found in the Go heap.
+// Prefer Store instead.
 //
 //go:nosplit
 func (u *UnsafePointer) StoreNoWB(value unsafe.Pointer) {
 	StorepNoWB(unsafe.Pointer(&u.value), value)
 }
 
+// Store updates the value atomically.
+func (u *UnsafePointer) Store(value unsafe.Pointer) {
+	storePointer(&u.value, value)
+}
+
+// provided by runtime
+//go:linkname storePointer
+func storePointer(ptr *unsafe.Pointer, new unsafe.Pointer)
+
 // CompareAndSwapNoWB atomically (with respect to other methods)
 // compares u's value with old, and if they're equal,
 // swaps u's value with new.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 // WARNING: As the name implies this operation does *not*
 // perform a write barrier on value, and so this operation may
 // hide pointers from the GC. Use with care and sparingly.
 // It is safe to use with values not found in the Go heap.
+// Prefer CompareAndSwap instead.
 //
 //go:nosplit
 func (u *UnsafePointer) CompareAndSwapNoWB(old, new unsafe.Pointer) bool {
 	return Casp1(&u.value, old, new)
 }
 
+// CompareAndSwap atomically compares u's value with old,
+// and if they're equal, swaps u's value with new.
+// It reports whether the swap ran.
+func (u *UnsafePointer) CompareAndSwap(old, new unsafe.Pointer) bool {
+	return casPointer(&u.value, old, new)
+}
+
+func casPointer(ptr *unsafe.Pointer, old, new unsafe.Pointer) bool
+
 // Pointer is an atomic pointer of type *T.
 type Pointer[T any] struct {
 	u UnsafePointer
@@ -518,28 +531,43 @@ func (p *Pointer[T]) Load() *T {
 // perform a write barrier on value, and so this operation may
 // hide pointers from the GC. Use with care and sparingly.
 // It is safe to use with values not found in the Go heap.
+// Prefer Store instead.
 //
 //go:nosplit
 func (p *Pointer[T]) StoreNoWB(value *T) {
 	p.u.StoreNoWB(unsafe.Pointer(value))
 }
 
+// Store updates the value atomically.
+//go:nosplit
+func (p *Pointer[T]) Store(value *T) {
+	p.u.Store(unsafe.Pointer(value))
+}
+
 // CompareAndSwapNoWB atomically (with respect to other methods)
 // compares u's value with old, and if they're equal,
 // swaps u's value with new.
-//
-// Returns true if the operation succeeded.
+// It reports whether the swap ran.
 //
 // WARNING: As the name implies this operation does *not*
 // perform a write barrier on value, and so this operation may
 // hide pointers from the GC. Use with care and sparingly.
 // It is safe to use with values not found in the Go heap.
+// Prefer CompareAndSwap instead.
 //
 //go:nosplit
 func (p *Pointer[T]) CompareAndSwapNoWB(old, new *T) bool {
 	return p.u.CompareAndSwapNoWB(unsafe.Pointer(old), unsafe.Pointer(new))
 }
 
+// CompareAndSwap atomically (with respect to other methods)
+// compares u's value with old, and if they're equal,
+// swaps u's value with new.
+// It reports whether the swap ran.
+func (p *Pointer[T]) CompareAndSwap(old, new *T) bool {
+	return p.u.CompareAndSwap(unsafe.Pointer(old), unsafe.Pointer(new))
+}
+
 // noCopy may be embedded into structs which must not be copied
 // after the first use.
 //