From ff5b4bae99c7dc8bb57660bb579d5df4ab31b1ce Mon Sep 17 00:00:00 2001
From: Joe Tsai <joetsai@digital-static.net>
Date: Wed, 18 Dec 2024 17:11:22 -0800
Subject: [PATCH] syncs: add MutexValue (#14422)

MutexValue is simply a value guarded by a mutex.
For any type that is not pointer-sized,
MutexValue will perform much better than AtomicValue
since it will not incur an allocation boxing the value
into an interface value (which is how Go's atomic.Value
is implemented under-the-hood).

Updates #cleanup

Signed-off-by: Joe Tsai <joetsai@digital-static.net>
---
 syncs/syncs.go      | 62 +++++++++++++++++++++++++++++++++++++++++++++
 syncs/syncs_test.go | 34 +++++++++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/syncs/syncs.go b/syncs/syncs.go
index acc0c88f2..337fca755 100644
--- a/syncs/syncs.go
+++ b/syncs/syncs.go
@@ -25,6 +25,7 @@ func initClosedChan() <-chan struct{} {
 }
 
 // AtomicValue is the generic version of [atomic.Value].
+// See [MutexValue] for guidance on whether to use this type.
 type AtomicValue[T any] struct {
 	v atomic.Value
 }
@@ -74,6 +75,67 @@ func (v *AtomicValue[T]) CompareAndSwap(oldV, newV T) (swapped bool) {
 	return v.v.CompareAndSwap(wrappedValue[T]{oldV}, wrappedValue[T]{newV})
 }
 
+// MutexValue is a value protected by a mutex.
+//
+// AtomicValue, [MutexValue], [atomic.Pointer] are similar and
+// overlap in their use cases.
+//
+//   - Use [atomic.Pointer] if the value being stored is a pointer and
+//     you only ever need load and store operations.
+//     An atomic pointer only occupies 1 word of memory.
+//
+//   - Use [MutexValue] if the value being stored is not a pointer or
+//     you need the ability for a mutex to protect a set of operations
+//     performed on the value.
+//     A mutex-guarded value occupies 1 word of memory plus
+//     the memory representation of T.
+//
+//   - AtomicValue is useful for non-pointer types that happen to
+//     have the memory layout of a single pointer.
+//     Examples include a map, channel, func, or a single field struct
+//     that contains any prior types.
+//     An atomic value occupies 2 words of memory.
+//     Consequently, Storing of non-pointer types always allocates.
+//
+// Note that [AtomicValue] has the ability to report whether it was set
+// while [MutexValue] lacks the ability to detect if the value was set
+// and it happens to be the zero value of T. If such a use case is
+// necessary, then you could consider wrapping T in [opt.Value].
+type MutexValue[T any] struct {
+	mu sync.Mutex
+	v  T
+}
+
+// WithLock calls f with a pointer to the value while holding the lock.
+// The provided pointer must not leak beyond the scope of the call.
+func (m *MutexValue[T]) WithLock(f func(p *T)) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	f(&m.v)
+}
+
+// Load returns a shallow copy of the underlying value.
+func (m *MutexValue[T]) Load() T {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.v
+}
+
+// Store stores a shallow copy of the provided value.
+func (m *MutexValue[T]) Store(v T) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.v = v
+}
+
+// Swap stores new into m and returns the previous value.
+func (m *MutexValue[T]) Swap(new T) (old T) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	old, m.v = m.v, new
+	return old
+}
+
 // WaitGroupChan is like a sync.WaitGroup, but has a chan that closes
 // on completion that you can wait on. (This, you can only use the
 // value once)
diff --git a/syncs/syncs_test.go b/syncs/syncs_test.go
index ee3711e76..901d42948 100644
--- a/syncs/syncs_test.go
+++ b/syncs/syncs_test.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 )
@@ -65,6 +66,39 @@ func TestAtomicValue(t *testing.T) {
 	}
 }
 
+func TestMutexValue(t *testing.T) {
+	var v MutexValue[time.Time]
+	if n := int(testing.AllocsPerRun(1000, func() {
+		v.Store(v.Load())
+		v.WithLock(func(*time.Time) {})
+	})); n != 0 {
+		t.Errorf("AllocsPerRun = %d, want 0", n)
+	}
+
+	now := time.Now()
+	v.Store(now)
+	if !v.Load().Equal(now) {
+		t.Errorf("Load = %v, want %v", v.Load(), now)
+	}
+
+	var group WaitGroup
+	var v2 MutexValue[int]
+	var sum int
+	for i := range 10 {
+		group.Go(func() {
+			old1 := v2.Load()
+			old2 := v2.Swap(old1 + i)
+			delta := old2 - old1
+			v2.WithLock(func(p *int) { *p += delta })
+		})
+		sum += i
+	}
+	group.Wait()
+	if v2.Load() != sum {
+		t.Errorf("Load = %v, want %v", v2.Load(), sum)
+	}
+}
+
 func TestWaitGroupChan(t *testing.T) {
 	wg := NewWaitGroupChan()