1 // Copyright 2019 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
5 // Package memoize supports memoizing the return values of functions with
6 // idempotent results that are expensive to compute.
8 // To use this package, build a store and use it to acquire handles with the
21 "golang.org/x/tools/internal/xcontext"
25 panicOnDestroyed = flag.Bool("memoize_panic_on_destroyed", false,
26 "Panic when a destroyed generation is read rather than returning an error. "+
27 "Panicking may make it easier to debug lifetime errors, especially when "+
28 "used with GOTRACEBACK=crash to see all running goroutines.")
31 // Store binds keys to functions, returning handles that can be used to access
32 // the functions results.
35 // handles is the set of values stored.
36 handles map[interface{}]*Handle
38 // generations is the set of generations live in this store.
39 generations map[*Generation]struct{}
42 // Generation creates a new Generation associated with s. Destroy must be
43 // called on the returned Generation once it is no longer in use. name is
44 // for debugging purposes only.
45 func (s *Store) Generation(name string) *Generation {
49 s.handles = map[interface{}]*Handle{}
50 s.generations = map[*Generation]struct{}{}
52 g := &Generation{store: s, name: name}
53 s.generations[g] = struct{}{}
57 // A Generation is a logical point in time of the cache life-cycle. Cache
58 // entries associated with a Generation will not be removed until the
59 // Generation is destroyed.
60 type Generation struct {
61 // destroyed is 1 after the generation is destroyed. Atomic.
65 // wg tracks the reference count of this generation.
69 // Destroy waits for all operations referencing g to complete, then removes
70 // all references to g from cache entries. Cache entries that no longer
71 // reference any non-destroyed generation are removed. Destroy must be called
72 // exactly once for each generation.
73 func (g *Generation) Destroy() {
75 atomic.StoreUint32(&g.destroyed, 1)
77 defer g.store.mu.Unlock()
78 for k, e := range g.store.handles {
80 if _, ok := e.generations[g]; ok {
81 delete(e.generations, g) // delete even if it's dead, in case of dangling references to the entry.
82 if len(e.generations) == 0 {
83 delete(g.store.handles, k)
84 e.state = stateDestroyed
85 if e.cleanup != nil && e.value != nil {
92 delete(g.store.generations, g)
95 // Acquire creates a new reference to g, and returns a func to release that
97 func (g *Generation) Acquire(ctx context.Context) func() {
98 destroyed := atomic.LoadUint32(&g.destroyed)
103 panic("acquire on destroyed generation " + g.name)
109 // Arg is a marker interface that can be embedded to indicate a type is
110 // intended for use as a Function argument.
111 type Arg interface{ memoizeArg() }
113 // Function is the type for functions that can be memoized.
114 // The result must be a pointer.
115 type Function func(ctx context.Context, arg Arg) interface{}
126 // Handle is returned from a store when a key is bound to a function.
127 // It is then used to access the results of that function.
129 // A Handle starts out in idle state, waiting for something to demand its
130 // evaluation. It then transitions into running state. While it's running,
131 // waiters tracks the number of Get calls waiting for a result, and the done
132 // channel is used to notify waiters of the next state transition. Once the
133 // evaluation finishes, value is set, state changes to completed, and done
134 // is closed, unblocking waiters. Alternatively, as Get calls are cancelled,
135 // they decrement waiters. If it drops to zero, the inner context is cancelled,
136 // computation is abandoned, and state resets to idle to start the process over
142 // generations is the set of generations in which this handle is valid.
143 generations map[*Generation]struct{}
146 // done is set in running state, and closed when exiting it.
148 // cancel is set in running state. It cancels computation.
149 cancel context.CancelFunc
150 // waiters is the number of Gets outstanding.
152 // the function that will be used to populate the value
154 // value is set in completed state.
156 // cleanup, if non-nil, is used to perform any necessary clean-up on values
157 // produced by function.
158 cleanup func(interface{})
161 // Bind returns a handle for the given key and function.
163 // Each call to bind will return the same handle if it is already bound. Bind
164 // will always return a valid handle, creating one if needed. Each key can
165 // only have one handle at any given time. The value will be held at least
166 // until the associated generation is destroyed. Bind does not cause the value
169 // If cleanup is non-nil, it will be called on any non-nil values produced by
170 // function when they are no longer referenced.
171 func (g *Generation) Bind(key interface{}, function Function, cleanup func(interface{})) *Handle {
172 // panic early if the function is nil
173 // it would panic later anyway, but in a way that was much harder to debug
175 panic("the function passed to bind must not be nil")
177 if atomic.LoadUint32(&g.destroyed) != 0 {
178 panic("operation on destroyed generation " + g.name)
181 defer g.store.mu.Unlock()
182 h, ok := g.store.handles[key]
187 generations: map[*Generation]struct{}{g: {}},
190 g.store.handles[key] = h
195 if _, ok := h.generations[g]; !ok {
196 h.generations[g] = struct{}{}
201 // Stats returns the number of each type of value in the store.
202 func (s *Store) Stats() map[reflect.Type]int {
206 result := map[reflect.Type]int{}
207 for k := range s.handles {
208 result[reflect.TypeOf(k)]++
213 // DebugOnlyIterate iterates through all live cache entries and calls f on them.
214 // It should only be used for debugging purposes.
215 func (s *Store) DebugOnlyIterate(f func(k, v interface{})) {
219 for k, e := range s.handles {
222 if e.state == stateCompleted {
233 func (g *Generation) Inherit(hs ...*Handle) {
234 for _, h := range hs {
235 if atomic.LoadUint32(&g.destroyed) != 0 {
236 panic("inherit on destroyed generation " + g.name)
241 if h.state == stateDestroyed {
242 panic(fmt.Sprintf("inheriting destroyed handle %#v (type %T) into generation %v", h.key, h.key, g.name))
244 h.generations[g] = struct{}{}
248 // Cached returns the value associated with a handle.
250 // It will never cause the value to be generated.
251 // It will return the cached value, if present.
252 func (h *Handle) Cached(g *Generation) interface{} {
255 if _, ok := h.generations[g]; !ok {
258 if h.state == stateCompleted {
264 // Get returns the value associated with a handle.
266 // If the value is not yet ready, the underlying function will be invoked.
267 // If ctx is cancelled, Get returns nil.
268 func (h *Handle) Get(ctx context.Context, g *Generation, arg Arg) (interface{}, error) {
269 release := g.Acquire(ctx)
272 if ctx.Err() != nil {
273 return nil, ctx.Err()
276 if _, ok := h.generations[g]; !ok {
279 err := fmt.Errorf("reading key %#v: generation %v is not known", h.key, g.name)
280 if *panicOnDestroyed && ctx.Err() != nil {
287 return h.run(ctx, g, arg)
295 err := fmt.Errorf("Get on destroyed entry %#v (type %T) in generation %v", h.key, h.key, g.name)
296 if *panicOnDestroyed {
301 panic("unknown state")
305 // run starts h.function and returns the result. h.mu must be locked.
306 func (h *Handle) run(ctx context.Context, g *Generation, arg Arg) (interface{}, error) {
307 childCtx, cancel := context.WithCancel(xcontext.Detach(ctx))
309 h.state = stateRunning
310 h.done = make(chan struct{})
311 function := h.function // Read under the lock
313 // Make sure that the generation isn't destroyed while we're running in it.
314 release := g.Acquire(ctx)
317 // Just in case the function does something expensive without checking
318 // the context, double-check we're still alive.
319 if childCtx.Err() != nil {
322 v := function(childCtx, arg)
323 if childCtx.Err() != nil {
324 // It's possible that v was computed despite the context cancellation. In
325 // this case we should ensure that it is cleaned up.
326 if h.cleanup != nil && v != nil {
334 // It's theoretically possible that the handle has been cancelled out
335 // of the run that started us, and then started running again since we
336 // checked childCtx above. Even so, that should be harmless, since each
337 // run should produce the same results.
338 if h.state != stateRunning {
339 // v will never be used, so ensure that it is cleaned up.
340 if h.cleanup != nil && v != nil {
345 // At this point v will be cleaned up whenever h is destroyed.
348 h.state = stateCompleted
355 // wait waits for the value to be computed, or ctx to be cancelled. h.mu must be locked.
356 func (h *Handle) wait(ctx context.Context) (interface{}, error) {
365 if h.state == stateCompleted {
373 if h.waiters == 0 && h.state == stateRunning {
380 return nil, ctx.Err()