From d209b032abbc5f86301d296e1919abdc302e3cfc Mon Sep 17 00:00:00 2001
From: Joe Tsai <joetsai@digital-static.net>
Date: Thu, 11 Jul 2024 16:16:30 -0700
Subject: [PATCH] syncs: add Map.WithLock to allow mutations to the underlying
 map (#8101)

Some operations cannot be implemented with the prior API:
* Iterating over the map and deleting keys
* Iterating over the map and replacing items
* Calling APIs that expect a native Go map

Add a Map.WithLock method that acquires a write-lock on the map
and then calls a user-provided closure with the underlying Go map.
This allows users to interact with the Map as a regular Go map,
but with the gaurantees that it is concurrent safe.

Updates tailscale/corp#9115

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

diff --git a/syncs/syncs.go b/syncs/syncs.go
index c3f729a90..0d40204d2 100644
--- a/syncs/syncs.go
+++ b/syncs/syncs.go
@@ -252,8 +252,10 @@ func (m *Map[K, V]) Delete(key K) {
 	delete(m.m, key)
 }
 
-// Range iterates over the map in undefined order calling f for each entry.
+// Range iterates over the map in an undefined order calling f for each entry.
 // Iteration stops if f returns false. Map changes are blocked during iteration.
+// A read lock is held for the entire duration of the iteration.
+// Use the [WithLock] method instead to mutate the map during iteration.
 func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 	m.mu.RLock()
 	defer m.mu.RUnlock()
@@ -264,6 +266,15 @@ func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 	}
 }
 
+// WithLock calls f with the underlying map.
+// Use of m2 must not escape the duration of this call.
+// The write-lock is held for the entire duration of this call.
+func (m *Map[K, V]) WithLock(f func(m2 map[K]V)) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	f(m.m)
+}
+
 // Len returns the length of the map.
 func (m *Map[K, V]) Len() int {
 	m.mu.RLock()
diff --git a/syncs/syncs_test.go b/syncs/syncs_test.go
index 424d51794..0748dcb72 100644
--- a/syncs/syncs_test.go
+++ b/syncs/syncs_test.go
@@ -7,7 +7,6 @@ import (
 	"context"
 	"io"
 	"os"
-	"sync"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -189,19 +188,11 @@ func TestMap(t *testing.T) {
 
 	t.Run("LoadOrStore", func(t *testing.T) {
 		var m Map[string, string]
-		var wg sync.WaitGroup
-		wg.Add(2)
+		var wg WaitGroup
 		var ok1, ok2 bool
-		go func() {
-			defer wg.Done()
-			_, ok1 = m.LoadOrStore("", "")
-		}()
-		go func() {
-			defer wg.Done()
-			_, ok2 = m.LoadOrStore("", "")
-		}()
+		wg.Go(func() { _, ok1 = m.LoadOrStore("", "") })
+		wg.Go(func() { _, ok2 = m.LoadOrStore("", "") })
 		wg.Wait()
-
 		if ok1 == ok2 {
 			t.Errorf("exactly one LoadOrStore should load")
 		}