From 700c7b95ae001244e60e820dcc4f63ae4f5fc5b1 Mon Sep 17 00:00:00 2001 From: Michael Anthony Knyszek Date: Tue, 13 Aug 2024 16:31:27 +0000 Subject: [PATCH] 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 Auto-Submit: Michael Knyszek Reviewed-by: David Chase --- src/internal/sync/hashtriemap.go | 73 +++++++++++ src/internal/sync/hashtriemap_test.go | 170 +++++++++++++++++++++++++- 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 { -- 2.48.1