internal/sync: add LoadAndDelete to HashTrieMap

This change adds the LoadAndDelete operation (with the same semantics as
sync.Map's LoadAndDelete) to HashTrieMap.

Change-Id: Id6777dffcd3ebc98490aa51f0e85e59a56f63074
Reviewed-on: https://go-review.googlesource.com/c/go/+/606456
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Michael Knyszek <mknyszek@google.com>
Reviewed-by: David Chase <drchase@google.com>

2 files changed, 238 insertions, 5 deletions

diff --git a/src/internal/sync/hashtriemap.go b/src/internal/sync/hashtriemap.go
index 6e66bc81d3..5862962e9b 100644
--- a/src/internal/sync/hashtriemap.go
+++ b/src/internal/sync/hashtriemap.go
@@ -304,6 +304,57 @@ func (ht *HashTrieMap[K, V]) CompareAndSwap(key K, old, new V) (swapped bool) {
 	return true
 }
 
+// LoadAndDelete deletes the value for a key, returning the previous value if any.
+// The loaded result reports whether the key was present.
+func (ht *HashTrieMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
+	ht.init()
+	hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
+
+	// Find a node with the key and compare with it. n != nil if we found the node.
+	i, hashShift, slot, n := ht.find(key, hash, nil, *new(V))
+	if n == nil {
+		if i != nil {
+			i.mu.Unlock()
+		}
+		return *new(V), false
+	}
+
+	// Try to delete the entry.
+	v, e, loaded := n.entry().loadAndDelete(key)
+	if !loaded {
+		// Nothing was actually deleted, which means the node is no longer there.
+		i.mu.Unlock()
+		return *new(V), false
+	}
+	if e != nil {
+		// We didn't actually delete the whole entry, just one entry in the chain.
+		// Nothing else to do, since the parent is definitely not empty.
+		slot.Store(&e.node)
+		i.mu.Unlock()
+		return v, true
+	}
+	// Delete the entry.
+	slot.Store(nil)
+
+	// Check if the node is now empty (and isn't the root), and delete it if able.
+	for i.parent != nil && i.empty() {
+		if hashShift == 8*goarch.PtrSize {
+			panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
+		}
+		hashShift += nChildrenLog2
+
+		// Delete the current node in the parent.
+		parent := i.parent
+		parent.mu.Lock()
+		i.dead.Store(true)
+		parent.children[(hash>>hashShift)&nChildrenMask].Store(nil)
+		i.mu.Unlock()
+		i = parent
+	}
+	i.mu.Unlock()
+	return v, true
+}
+
 // CompareAndDelete deletes the entry for key if its value is equal to old.
 // The value type must be comparable, otherwise this CompareAndDelete will panic.
 //
@@ -575,6 +626,28 @@ func (head *entry[K, V]) compareAndSwap(key K, old, new V, valEqual equalFunc) (
 	return head, false
 }
 
+// loadAndDelete deletes an entry in the overflow chain by key. Returns the value for the key, the new
+// entry chain and whether or not anything was loaded (and deleted).
+//
+// loadAndDelete must be called under the mutex of the indirect node which e is a child of.
+func (head *entry[K, V]) loadAndDelete(key K) (V, *entry[K, V], bool) {
+	if head.key == key {
+		// Drop the head of the list.
+		return head.value, head.overflow.Load(), true
+	}
+	i := &head.overflow
+	e := i.Load()
+	for e != nil {
+		if e.key == key {
+			i.Store(e.overflow.Load())
+			return e.value, head, true
+		}
+		i = &e.overflow
+		e = e.overflow.Load()
+	}
+	return *new(V), head, false
+}
+
 // compareAndDelete deletes an entry in the overflow chain if both the key and value compare
 // equal. Returns the new entry chain and whether or not anything was deleted.
 //
diff --git a/src/internal/sync/hashtriemap_test.go b/src/internal/sync/hashtriemap_test.go
index f3e1073c4f..cca7512350 100644
--- a/src/internal/sync/hashtriemap_test.go
+++ b/src/internal/sync/hashtriemap_test.go
@@ -141,7 +141,7 @@ func testHashTrieMap(t *testing.T, newMap func() *isync.HashTrieMap[string, int]
 			expectMissing(t, s, 0)(m.Load(s))
 		}
 	})
-	t.Run("ConcurrentLifecycleUnsharedKeys", func(t *testing.T) {
+	t.Run("ConcurrentCompareAndDeleteUnsharedKeys", func(t *testing.T) {
 		m := newMap()
 
 		gmp := runtime.GOMAXPROCS(-1)
@@ -266,7 +266,7 @@ func testHashTrieMap(t *testing.T, newMap func() *isync.HashTrieMap[string, int]
 			}
 		}
 	})
-	t.Run("ConcurrentLifecycleCompareAndSwapUnsharedKeys", func(t *testing.T) {
+	t.Run("ConcurrentCompareAndSwapUnsharedKeys", func(t *testing.T) {
 		m := newMap()
 
 		gmp := runtime.GOMAXPROCS(-1)
@@ -300,7 +300,7 @@ func testHashTrieMap(t *testing.T, newMap func() *isync.HashTrieMap[string, int]
 		}
 		wg.Wait()
 	})
-	t.Run("ConcurrentLifecycleCompareAndSwapAndDeleteUnsharedKeys", func(t *testing.T) {
+	t.Run("ConcurrentCompareAndSwapAndDeleteUnsharedKeys", func(t *testing.T) {
 		m := newMap()
 
 		gmp := runtime.GOMAXPROCS(-1)
@@ -423,7 +423,7 @@ func testHashTrieMap(t *testing.T, newMap func() *isync.HashTrieMap[string, int]
 			}
 		}
 	})
-	t.Run("ConcurrentLifecycleSwapUnsharedKeys", func(t *testing.T) {
+	t.Run("ConcurrentSwapUnsharedKeys", func(t *testing.T) {
 		m := newMap()
 
 		gmp := runtime.GOMAXPROCS(-1)
@@ -457,7 +457,7 @@ func testHashTrieMap(t *testing.T, newMap func() *isync.HashTrieMap[string, int]
 		}
 		wg.Wait()
 	})
-	t.Run("ConcurrentLifecycleSwapAndDeleteUnsharedKeys", func(t *testing.T) {
+	t.Run("ConcurrentSwapAndDeleteUnsharedKeys", func(t *testing.T) {
 		m := newMap()
 
 		gmp := runtime.GOMAXPROCS(-1)
@@ -520,6 +520,142 @@ func testHashTrieMap(t *testing.T, newMap func() *isync.HashTrieMap[string, int]
 		}
 		wg.Wait()
 	})
+	t.Run("LoadAndDeleteAll", func(t *testing.T) {
+		m := newMap()
+
+		for range 3 {
+			for i, s := range testData {
+				expectMissing(t, s, 0)(m.Load(s))
+				expectStored(t, s, i)(m.LoadOrStore(s, i))
+				expectPresent(t, s, i)(m.Load(s))
+				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
+			}
+			for i, s := range testData {
+				expectPresent(t, s, i)(m.Load(s))
+				expectLoadedFromDelete(t, s, i)(m.LoadAndDelete(s))
+				expectMissing(t, s, 0)(m.Load(s))
+				expectNotLoadedFromDelete(t, s, 0)(m.LoadAndDelete(s))
+			}
+			for _, s := range testData {
+				expectMissing(t, s, 0)(m.Load(s))
+			}
+		}
+	})
+	t.Run("LoadAndDeleteOne", func(t *testing.T) {
+		m := newMap()
+
+		for i, s := range testData {
+			expectMissing(t, s, 0)(m.Load(s))
+			expectStored(t, s, i)(m.LoadOrStore(s, i))
+			expectPresent(t, s, i)(m.Load(s))
+			expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
+		}
+		expectPresent(t, testData[15], 15)(m.Load(testData[15]))
+		expectLoadedFromDelete(t, testData[15], 15)(m.LoadAndDelete(testData[15]))
+		expectMissing(t, testData[15], 0)(m.Load(testData[15]))
+		expectNotLoadedFromDelete(t, testData[15], 0)(m.LoadAndDelete(testData[15]))
+		for i, s := range testData {
+			if i == 15 {
+				expectMissing(t, s, 0)(m.Load(s))
+			} else {
+				expectPresent(t, s, i)(m.Load(s))
+			}
+		}
+	})
+	t.Run("LoadAndDeleteMultiple", func(t *testing.T) {
+		m := newMap()
+
+		for i, s := range testData {
+			expectMissing(t, s, 0)(m.Load(s))
+			expectStored(t, s, i)(m.LoadOrStore(s, i))
+			expectPresent(t, s, i)(m.Load(s))
+			expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
+		}
+		for _, i := range []int{1, 105, 6, 85} {
+			expectPresent(t, testData[i], i)(m.Load(testData[i]))
+			expectLoadedFromDelete(t, testData[i], i)(m.LoadAndDelete(testData[i]))
+			expectMissing(t, testData[i], 0)(m.Load(testData[i]))
+			expectNotLoadedFromDelete(t, testData[i], 0)(m.LoadAndDelete(testData[i]))
+		}
+		for i, s := range testData {
+			if i == 1 || i == 105 || i == 6 || i == 85 {
+				expectMissing(t, s, 0)(m.Load(s))
+			} else {
+				expectPresent(t, s, i)(m.Load(s))
+			}
+		}
+	})
+	t.Run("AllLoadAndDelete", func(t *testing.T) {
+		m := newMap()
+
+		testAll(t, m, testDataMap(testData[:]), func(s string, i int) bool {
+			expectLoadedFromDelete(t, s, i)(m.LoadAndDelete(s))
+			return true
+		})
+		for _, s := range testData {
+			expectMissing(t, s, 0)(m.Load(s))
+		}
+	})
+	t.Run("ConcurrentLoadAndDeleteUnsharedKeys", func(t *testing.T) {
+		m := newMap()
+
+		gmp := runtime.GOMAXPROCS(-1)
+		var wg sync.WaitGroup
+		for i := range gmp {
+			wg.Add(1)
+			go func(id int) {
+				defer wg.Done()
+
+				makeKey := func(s string) string {
+					return s + "-" + strconv.Itoa(id)
+				}
+				for _, s := range testData {
+					key := makeKey(s)
+					expectMissing(t, key, 0)(m.Load(key))
+					expectStored(t, key, id)(m.LoadOrStore(key, id))
+					expectPresent(t, key, id)(m.Load(key))
+					expectLoaded(t, key, id)(m.LoadOrStore(key, 0))
+				}
+				for _, s := range testData {
+					key := makeKey(s)
+					expectPresent(t, key, id)(m.Load(key))
+					expectLoadedFromDelete(t, key, id)(m.LoadAndDelete(key))
+					expectMissing(t, key, 0)(m.Load(key))
+				}
+				for _, s := range testData {
+					key := makeKey(s)
+					expectMissing(t, key, 0)(m.Load(key))
+				}
+			}(i)
+		}
+		wg.Wait()
+	})
+	t.Run("ConcurrentLoadAndDeleteSharedKeys", func(t *testing.T) {
+		m := newMap()
+
+		// Load up the map.
+		for i, s := range testData {
+			expectMissing(t, s, 0)(m.Load(s))
+			expectStored(t, s, i)(m.LoadOrStore(s, i))
+		}
+		gmp := runtime.GOMAXPROCS(-1)
+		var wg sync.WaitGroup
+		for i := range gmp {
+			wg.Add(1)
+			go func(id int) {
+				defer wg.Done()
+
+				for _, s := range testData {
+					m.LoadAndDelete(s)
+					expectMissing(t, s, 0)(m.Load(s))
+				}
+				for _, s := range testData {
+					expectMissing(t, s, 0)(m.Load(s))
+				}
+			}(i)
+		}
+		wg.Wait()
+	})
 }
 
 func testAll[K, V comparable](t *testing.T, m *isync.HashTrieMap[K, V], testData map[K]V, yield func(K, V) bool) {
@@ -676,6 +812,30 @@ func expectNotLoadedFromSwap[K, V comparable](t *testing.T, key K, new V) func(o
 	}
 }
 
+func expectLoadedFromDelete[K, V comparable](t *testing.T, key K, want V) func(got V, loaded bool) {
+	t.Helper()
+	return func(got V, loaded bool) {
+		t.Helper()
+
+		if !loaded {
+			t.Errorf("expected key %v to be in map to be deleted", key)
+		} else if want != got {
+			t.Errorf("key %v was deleted with value %v, but expected it to have value %v", key, got, want)
+		}
+	}
+}
+
+func expectNotLoadedFromDelete[K, V comparable](t *testing.T, key K, _ V) func(old V, loaded bool) {
+	t.Helper()
+	return func(old V, loaded bool) {
+		t.Helper()
+
+		if loaded {
+			t.Errorf("expected key %v to not be in map, but found value %v for it", key, old)
+		}
+	}
+}
+
 func testDataMap(data []string) map[string]int {
 	m := make(map[string]int)
 	for i, s := range data {