1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
|
// Copyright (C) 2023 Luke Shumaker <lukeshu@lukeshu.com>
//
// SPDX-License-Identifier: GPL-2.0-or-later
package syncutil
import (
"context"
"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
c chan struct{}
}
// 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 = &MapOnceVal[V]{
c: make(chan struct{}),
}
}
_actual, loaded := m.Inner.LoadOrStore(key, _value)
if loaded {
m.pool.Put(_value)
<-_actual.c
} else {
_actual.V = do(key)
close(_actual.c)
}
return _actual.V, loaded
}
// TryLoadOrDo is like LoadOrDo, but obeys context cancellation. If a
// call is cancelled, the call to "do" continues running in a separate
// goroutine, in case other LoadOrDo calls are waiting on it. If a
// call is cancelled, the error from ctx.Err() is returned, otherwise
// err is nil.
func (m *MapOnce[K, V, M]) TryLoadOrDo(ctx context.Context, key K, do func(K) V) (actual V, loaded bool, err error) {
_value, _ := m.pool.Get()
if _value == nil {
_value = &MapOnceVal[V]{
c: make(chan struct{}),
}
}
_actual, loaded := m.Inner.LoadOrStore(key, _value)
if loaded {
m.pool.Put(_value)
} else {
go func() {
_actual.V = do(key)
close(_actual.c)
}()
}
select {
case <-ctx.Done():
var zero V
return zero, false, ctx.Err() //nolint:wrapcheck // We're too low level for that to be useful.
case <-_actual.c:
return _actual.V, loaded, nil
}
}
|
