cachemap: Revert the previous change, use a Pool to avoid allocations

1 files changed, 26 insertions, 17 deletions

diff --git a/cachemap.go b/cachemap.go
index bc6a70a..3294aa2 100644
--- a/cachemap.go
+++ b/cachemap.go
@@ -10,10 +10,7 @@ import (
 
 type cacheVal[V any] struct {
 	wg sync.WaitGroup
-	// 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
+	v  V
 }
 
 // The techniques used by CacheMap are similar to the techniques used
@@ -28,6 +25,7 @@ type cacheVal[V any] struct {
 // and more time overhead.
 type CacheMap[K mapkey, V any] struct {
 	inner Map[K, *cacheVal[V]]
+	pool  Pool[*cacheVal[V]]
 }
 
 func (m *CacheMap[K, V]) Delete(key K) {
@@ -44,7 +42,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
@@ -59,7 +57,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
@@ -68,16 +66,18 @@ 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) {
-	_value := &cacheVal[V]{}
-	_value.wg.Add(1)
+	_value, _ := m.pool.Get()
+	if _value == nil {
+		_value = new(cacheVal[V])
+	}
+	_value.v = value
 	_actual, loaded := m.inner.LoadOrStore(key, _value)
 	if loaded {
+		*_value = cacheVal[V]{}
+		m.pool.Put(_value)
 		_actual.wg.Wait()
-	} else {
-		_actual.v = &value
-		_actual.wg.Done()
 	}
-	return *_actual.v, loaded
+	return _actual.v, loaded
 }
 
 // LoadOrCompute returns the existing value for the key if present.
@@ -87,19 +87,28 @@ func (m *CacheMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 // value for that key, a latter call blocks until the computation is
 // complete, and then returns the initial call's value.
 func (m *CacheMap[K, V]) LoadOrCompute(key K, fn func(K) V) (actual V, loaded bool) {
-	_value := &cacheVal[V]{}
+	_value, _ := m.pool.Get()
+	if _value == nil {
+		_value = new(cacheVal[V])
+	}
 	_value.wg.Add(1)
 	_actual, loaded := m.inner.LoadOrStore(key, _value)
 	if loaded {
+		*_value = cacheVal[V]{}
+		m.pool.Put(_value)
 		_actual.wg.Wait()
 	} else {
-		v := fn(key)
-		_actual.v = &v
+		_actual.v = fn(key)
 		_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})
+	_value, _ := m.pool.Get()
+	if _value == nil {
+		_value = new(cacheVal[V])
+	}
+	_value.v = value
+	m.inner.Store(key, _value)
 }