cachemap: Try to avoid unnecessary large allocations

1 files changed, 20 insertions, 9 deletions

diff --git a/cachemap.go b/cachemap.go
index 90d2274..bc6a70a 100644
--- a/cachemap.go
+++ b/cachemap.go
@@ -10,7 +10,10 @@ import (
 
 type cacheVal[V any] struct {
 	wg sync.WaitGroup
-	v  V
+	// This is a pointer so that LoarOrStore and LoadOrCompute
+	// don't doesn't need to allocate a (potentially large) `V` on
+	// the heap unless they're actually storing it.
+	v *V
 }
 
 // The techniques used by CacheMap are similar to the techniques used
@@ -41,7 +44,7 @@ func (m *CacheMap[K, V]) Load(key K) (value V, ok bool) {
 		return zero, false
 	}
 	_value.wg.Wait()
-	return _value.v, true
+	return *_value.v, true
 }
 
 // LoadAndDelete deletes the value for a key, returning the previous
@@ -56,7 +59,7 @@ func (m *CacheMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
 		return zero, false
 	}
 	_value.wg.Wait()
-	return _value.v, true
+	return *_value.v, true
 }
 
 // LoadOrStore returns the existing value for the key if
@@ -65,9 +68,16 @@ func (m *CacheMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
 // the value for that key is actively being computed by LoadOrCompute,
 // this blocks until the value has been computed.
 func (m *CacheMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
-	_actual, loaded := m.inner.LoadOrStore(key, &cacheVal[V]{v: value})
-	_actual.wg.Wait()
-	return _actual.v, loaded
+	_value := &cacheVal[V]{}
+	_value.wg.Add(1)
+	_actual, loaded := m.inner.LoadOrStore(key, _value)
+	if loaded {
+		_actual.wg.Wait()
+	} else {
+		_actual.v = &value
+		_actual.wg.Done()
+	}
+	return *_actual.v, loaded
 }
 
 // LoadOrCompute returns the existing value for the key if present.
@@ -83,12 +93,13 @@ func (m *CacheMap[K, V]) LoadOrCompute(key K, fn func(K) V) (actual V, loaded bo
 	if loaded {
 		_actual.wg.Wait()
 	} else {
-		_actual.v = fn(key)
+		v := fn(key)
+		_actual.v = &v
 		_actual.wg.Done()
 	}
-	return _actual.v, loaded
+	return *_actual.v, loaded
 }
 
 func (m *CacheMap[K, V]) Store(key K, value V) {
-	m.inner.Store(key, &cacheVal[V]{v: value})
+	m.inner.Store(key, &cacheVal[V]{v: &value})
 }