syncs: add Map.LoadOrInit for lazily initialized values

I was reviewing some code that was performing this by hand, and wanted
to suggest using syncs.Map, however as the code in question was
allocating a non-trivial structure this would be necessary to meet the
target.

Updates #cleanup

Signed-off-by: James Tucker <james@tailscale.com>
This commit is contained in:
James Tucker 2023-10-06 10:24:21 -07:00 committed by James Tucker
parent e4cb83b18b
commit 498f7ec663
2 changed files with 25 additions and 1 deletions

View File

@ -192,6 +192,26 @@ func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
return actual, loaded
}
// LoadOrInit returns the value for the given key if it exists
// otherwise f is called to construct the value to be set.
// The lock is held for the duration to prevent duplicate initialization.
func (m *Map[K, V]) LoadOrInit(key K, f func() V) (actual V, loaded bool) {
if actual, loaded := m.Load(key); loaded {
return actual, loaded
}
m.mu.Lock()
defer m.mu.Unlock()
if actual, loaded = m.m[key]; loaded {
return actual, loaded
}
loaded = false
actual = f()
mak.Set(&m.m, key, actual)
return actual, loaded
}
func (m *Map[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
m.mu.Lock()
defer m.mu.Unlock()

View File

@ -91,8 +91,11 @@ func TestMap(t *testing.T) {
if v, ok := m.LoadOrStore("two", 2); v != 2 || ok {
t.Errorf(`LoadOrStore("two", 2) = (%v, %v), want (2, false)`, v, ok)
}
if v, ok := m.LoadOrInit("three", func() int { return 3 }); v != 3 || ok {
t.Errorf(`LoadOrInit("three", 3) = (%v, %v), want (3, true)`, v, ok)
}
got := map[string]int{}
want := map[string]int{"one": 1, "two": 2}
want := map[string]int{"one": 1, "two": 2, "three": 3}
m.Range(func(k string, v int) bool {
got[k] = v
return true
@ -106,6 +109,7 @@ func TestMap(t *testing.T) {
if v, ok := m.LoadAndDelete("two"); v != 0 || ok {
t.Errorf(`LoadAndDelete("two) = (%v, %v), want (0, false)`, v, ok)
}
m.Delete("three")
m.Delete("one")
m.Delete("noexist")
got = map[string]int{}