diff --git a/syncs/syncs.go b/syncs/syncs.go index abfaba5d8..274e83b65 100644 --- a/syncs/syncs.go +++ b/syncs/syncs.go @@ -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() diff --git a/syncs/syncs_test.go b/syncs/syncs_test.go index bb289862b..b061aebaf 100644 --- a/syncs/syncs_test.go +++ b/syncs/syncs_test.go @@ -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{}