internal/sync: factor out lookup for CompareAndDelete in HashTrieMap

This lookup will be reused for other operations, like CompareAndSwap.

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

diff --git a/src/internal/sync/hashtriemap.go b/src/internal/sync/hashtriemap.go
index a7e833ef3777f482c7f614c7a6f0ecf63f47fce1..e2d3766d692ad7151a89f0c6d5de6425560d8db1 100644 (file)
--- a/src/internal/sync/hashtriemap.go
+++ b/src/internal/sync/hashtriemap.go
@@ -181,57 +181,16 @@ func (ht *HashTrieMap[K, V]) expand(oldEntry, newEntry *entry[K, V], newHash uin
 // (even if the old value is the nil interface value).
 func (ht *HashTrieMap[K, V]) CompareAndDelete(key K, old V) (deleted bool) {
        hash := ht.keyHash(abi.NoEscape(unsafe.Pointer(&key)), ht.seed)
-       var i *indirect[K, V]
-       var hashShift uint
-       var slot *atomic.Pointer[node[K, V]]
-       var n *node[K, V]
-       for {
-               // Find the key or return when there's nothing to delete.
-               i = ht.root
-               hashShift = 8 * goarch.PtrSize
-               found := false
-               for hashShift != 0 {
-                       hashShift -= nChildrenLog2
 
-                       slot = &i.children[(hash>>hashShift)&nChildrenMask]
-                       n = slot.Load()
-                       if n == nil {
-                               // Nothing to delete. Give up.
-                               return
-                       }
-                       if n.isEntry {
-                               // We found an entry. Check if it matches.
-                               if _, ok := n.entry().lookup(key, ht.keyEqual); !ok {
-                                       // No match, nothing to delete.
-                                       return
-                               }
-                               // We've got something to delete.
-                               found = true
-                               break
-                       }
-                       i = n.indirect()
-               }
-               if !found {
-                       panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
-               }
-
-               // Grab the lock and double-check what we saw.
-               i.mu.Lock()
-               n = slot.Load()
-               if !i.dead.Load() {
-                       if n == nil {
-                               // Valid node that doesn't contain what we need. Nothing to delete.
-                               i.mu.Unlock()
-                               return
-                       }
-                       if n.isEntry {
-                               // What we saw is still true, so we can continue with the delete.
-                               break
-                       }
+       // 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)
+       if n == nil {
+               if i != nil {
+                       i.mu.Unlock()
                }
-               // We have to start over.
-               i.mu.Unlock()
+               return false
        }
+
        // Try to delete the entry.
        e, deleted := n.entry().compareAndDelete(key, old, ht.keyEqual, ht.valEqual)
        if !deleted {
@@ -268,6 +227,58 @@ func (ht *HashTrieMap[K, V]) CompareAndDelete(key K, old V) (deleted bool) {
        return true
 }
 
+// compare searches the tree for a node that compares with key (hash must be the hash of key).
+//
+// Returns a non-nil node, which will always be an entry, if found.
+//
+// If i != nil then i.mu is locked, and it is the caller's responsibility to unlock it.
+func (ht *HashTrieMap[K, V]) find(key K, hash uintptr) (i *indirect[K, V], hashShift uint, slot *atomic.Pointer[node[K, V]], n *node[K, V]) {
+       for {
+               // Find the key or return when there's nothing to delete.
+               i = ht.root
+               hashShift = 8 * goarch.PtrSize
+               found := false
+               for hashShift != 0 {
+                       hashShift -= nChildrenLog2
+
+                       slot = &i.children[(hash>>hashShift)&nChildrenMask]
+                       n = slot.Load()
+                       if n == nil {
+                               // Nothing to compare with. Give up.
+                               i = nil
+                               return
+                       }
+                       if n.isEntry {
+                               // We found an entry. Check if it matches.
+                               if _, ok := n.entry().lookup(key, ht.keyEqual); !ok {
+                                       // No match, comparison failed.
+                                       i = nil
+                                       n = nil
+                                       return
+                               }
+                               // We've got a match. Prepare to perform an operation on the key.
+                               found = true
+                               break
+                       }
+                       i = n.indirect()
+               }
+               if !found {
+                       panic("internal/concurrent.HashMapTrie: ran out of hash bits while iterating")
+               }
+
+               // Grab the lock and double-check what we saw.
+               i.mu.Lock()
+               n = slot.Load()
+               if !i.dead.Load() && (n == nil || n.isEntry) {
+                       // Either we've got a valid node or the node is now nil under the lock.
+                       // In either case, we're done here.
+                       return
+               }
+               // We have to start over.
+               i.mu.Unlock()
+       }
+}
+
 // All returns an iter.Seq2 that produces all key-value pairs in the map.
 // The enumeration does not represent any consistent snapshot of the map,
 // but is guaranteed to visit each unique key-value pair only once. It is