syncutil: Rethink CacheMap in to MapOnce

2 files changed, 87 insertions, 114 deletions

diff --git a/syncutil/cachemap.go b/syncutil/cachemap.go
deleted file mode 100644
index 8c4dfc5..0000000
--- a/syncutil/cachemap.go
+++ /dev/null
@@ -1,114 +0,0 @@
-// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package syncutil
-
-import (
-	"git.lukeshu.com/go/containers/typedsync"
-)
-
-type cacheVal[V any] struct {
-	wg typedsync.WaitGroup
-	v  V
-}
-
-// The techniques used by CacheMap are similar to the techniques used
-// by encoding/json's type cache.
-
-// A CacheMap is similar to a typedsync.Map, but also has a
-// .LoadOrCompute sibling to .LoadOrStore.  This is useful for caching
-// the results of computation where it is undesirable for concurrent
-// calls to duplicate work.
-//
-// Compared to a plain typedsync.Map, a CacheMap has both more space
-// overhead and more time overhead.
-type CacheMap[K mapkey, V any] struct {
-	inner typedsync.Map[K, *cacheVal[V]]
-	pool  typedsync.Pool[*cacheVal[V]]
-}
-
-func (m *CacheMap[K, V]) Delete(key K) {
-	m.inner.Delete(key)
-}
-
-// Load returns the value stored in the map for a key.  If the value
-// for that key is actively being computed by LoadOrCompute, this
-// blocks until the value has been computed.
-func (m *CacheMap[K, V]) Load(key K) (value V, ok bool) {
-	_value, ok := m.inner.Load(key)
-	if !ok {
-		var zero V
-		return zero, false
-	}
-	_value.wg.Wait()
-	return _value.v, true
-}
-
-// LoadAndDelete deletes the value for a key, returning the previous
-// value if any.  The loaded result reports whether the key was
-// present.  If the value for that key was actively being computed by
-// LoadOrCompute, this immediately removes the partial value from the
-// CacheMap, but does not return until the computation has finished.
-func (m *CacheMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
-	_value, ok := m.inner.LoadAndDelete(key)
-	if !ok {
-		var zero V
-		return zero, false
-	}
-	_value.wg.Wait()
-	return _value.v, true
-}
-
-// LoadOrStore returns the existing value for the key if
-// present. Otherwise, it stores and returns the given value.  The
-// loaded result is true if the value was loaded, false if stored.  If
-// 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, _ := 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()
-	}
-	return _actual.v, loaded
-}
-
-// LoadOrCompute returns the existing value for the key if present.
-// Otherwise, it computes and stores a value using the given function.
-// The loaded result is true if the value was loaded, false if
-// computed.  If a prior call to LoadOrCompute is still computing the
-// 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, _ := 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 {
-		_actual.v = fn(key)
-		_actual.wg.Done()
-	}
-	return _actual.v, loaded
-}
-
-func (m *CacheMap[K, V]) Store(key K, value V) {
-	_value, _ := m.pool.Get()
-	if _value == nil {
-		_value = new(cacheVal[V])
-	}
-	_value.v = value
-	m.inner.Store(key, _value)
-}
diff --git a/syncutil/maponce.go b/syncutil/maponce.go
new file mode 100644
index 0000000..150c6ea
--- /dev/null
+++ b/syncutil/maponce.go
@@ -0,0 +1,87 @@
+// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package syncutil
+
+import (
+	"git.lukeshu.com/go/containers/typedsync"
+)
+
+// A MapOnce wraps a Map (typically a
+// [git.lukeshu.com/go/containers/typedsync.Map], but possibly other
+// backends), in order to provide a LoadOrDo method that allows
+// missing values to be constructed without duplicating
+// work.
+//
+// It is similar to a [golang.org/x/sync/singleflight.Group], but
+// values are persistent.
+//
+// It is similar to a map of [sync.Once] values, but without concerns
+// about initialization.
+//
+// [git.lukeshu.com/go/containers/typedsync.Map]: https://pkg.go.dev/git.lukeshu.com/go/containers/typedsync#Map
+// [golang.org/x/sync/singleflight.Group]: https://pkg.go.dev/golang.org/x/sync/singleflight#Group
+// [sync.Once]: https://pkg.go.dev/sync#Once
+type MapOnce[K mapkey, V any, M Map[K, *MapOnceVal[V]]] struct {
+	// The techniques used by MapOnce are similar to the
+	// techniques used by encoding/json's internal type cache.
+
+	Inner M
+
+	// Because LoadOrDo needs a "new" empty *MapOnceVal[V] that
+	// will be immediately discarded for "Load" operations, it is
+	// worth having a pool so that should-be-fast "Load"s don't
+	// trigger slow allocations and create GC pressure.
+	pool typedsync.Pool[*MapOnceVal[V]]
+}
+
+// A Map describes the parallel-safe "map" storage required by a
+// MapOnce.
+//
+// The canonical Map implementation is
+// git.lukeshu.com/go/containers/typedsync.Map.
+type Map[K mapkey, V any] interface {
+	Delete(K)
+	LoadOrStore(K, V) (actual V, loaded bool)
+}
+
+// A MapOnceVal is a values that MapOnce stores in to its underlying
+// Map.
+type MapOnceVal[V any] struct {
+	V  V
+	wg typedsync.WaitGroup
+}
+
+// Delete removes the value for a key.  If the value for that key is
+// actively being constructed by LoadOrDo, this immediately removes
+// the partial value from the underlying map, but outstanding LoadOrDo
+// calls will still behave as if Delete had not been called.
+func (m *MapOnce[K, V, M]) Delete(key K) {
+	m.Inner.Delete(key)
+}
+
+// LoadOrDo returns the existing value stored for a key, if present.
+// If not present, it calls the "do" function to construct the value,
+// then stores and returns that value.  The "loaded" result is true if
+// the value was loaded, false if constructed.  If a prior call to
+// LoadOrDo is still constructing the value for that key, a latter
+// call blocks until the constructing is complete, and then returns
+// the initial call's value.
+func (m *MapOnce[K, V, M]) LoadOrDo(key K, do func(K) V) (actual V, loaded bool) {
+	_value, _ := m.pool.Get()
+	if _value == nil {
+		_value = new(MapOnceVal[V])
+	}
+	_value.wg.Add(1)
+	_actual, loaded := m.Inner.LoadOrStore(key, _value)
+	if loaded {
+		*_value = MapOnceVal[V]{}
+		m.pool.Put(_value)
+		_actual.wg.Wait()
+	} else {
+		_actual.V = do(key)
+		_actual.wg.Done()
+	}
+	return _actual.V, loaded
+}