From 498f7ec6636c7177a0dcf900aed6d903bf0da29e Mon Sep 17 00:00:00 2001
From: James Tucker <james@tailscale.com>
Date: Fri, 6 Oct 2023 10:24:21 -0700
Subject: [PATCH] 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>
---
 syncs/syncs.go      | 20 ++++++++++++++++++++
 syncs/syncs_test.go |  6 +++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

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{}