whatcanGOwrong
This commit is contained in:
@@ -0,0 +1,335 @@
|
||||
// Copyright 2019 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package memoize defines a "promise" abstraction that enables
|
||||
// memoization of the result of calling an expensive but idempotent
|
||||
// function.
|
||||
//
|
||||
// Call p = NewPromise(f) to obtain a promise for the future result of
|
||||
// calling f(), and call p.Get() to obtain that result. All calls to
|
||||
// p.Get return the result of a single call of f().
|
||||
// Get blocks if the function has not finished (or started).
|
||||
//
|
||||
// A Store is a map of arbitrary keys to promises. Use Store.Promise
|
||||
// to create a promise in the store. All calls to Handle(k) return the
|
||||
// same promise as long as it is in the store. These promises are
|
||||
// reference-counted and must be explicitly released. Once the last
|
||||
// reference is released, the promise is removed from the store.
|
||||
package memoize
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"runtime/trace"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"golang.org/x/tools/internal/xcontext"
|
||||
)
|
||||
|
||||
// Function is the type of a function that can be memoized.
|
||||
//
|
||||
// If the arg is a RefCounted, its Acquire/Release operations are called.
|
||||
//
|
||||
// The argument must not materially affect the result of the function
|
||||
// in ways that are not captured by the promise's key, since if
|
||||
// Promise.Get is called twice concurrently, with the same (implicit)
|
||||
// key but different arguments, the Function is called only once but
|
||||
// its result must be suitable for both callers.
|
||||
//
|
||||
// The main purpose of the argument is to avoid the Function closure
|
||||
// needing to retain large objects (in practice: the snapshot) in
|
||||
// memory that can be supplied at call time by any caller.
|
||||
type Function func(ctx context.Context, arg interface{}) interface{}
|
||||
|
||||
// A RefCounted is a value whose functional lifetime is determined by
|
||||
// reference counting.
|
||||
//
|
||||
// Its Acquire method is called before the Function is invoked, and
|
||||
// the corresponding release is called when the Function returns.
|
||||
// Usually both events happen within a single call to Get, so Get
|
||||
// would be fine with a "borrowed" reference, but if the context is
|
||||
// cancelled, Get may return before the Function is complete, causing
|
||||
// the argument to escape, and potential premature destruction of the
|
||||
// value. For a reference-counted type, this requires a pair of
|
||||
// increment/decrement operations to extend its life.
|
||||
type RefCounted interface {
|
||||
// Acquire prevents the value from being destroyed until the
|
||||
// returned function is called.
|
||||
Acquire() func()
|
||||
}
|
||||
|
||||
// A Promise represents the future result of a call to a function.
|
||||
type Promise struct {
|
||||
debug string // for observability
|
||||
|
||||
// refcount is the reference count in the containing Store, used by
|
||||
// Store.Promise. It is guarded by Store.promisesMu on the containing Store.
|
||||
refcount int32
|
||||
|
||||
mu sync.Mutex
|
||||
|
||||
// A Promise starts out IDLE, waiting for something to demand
|
||||
// its evaluation. It then transitions into RUNNING state.
|
||||
//
|
||||
// While RUNNING, waiters tracks the number of Get calls
|
||||
// waiting for a result, and the done channel is used to
|
||||
// notify waiters of the next state transition. Once
|
||||
// evaluation finishes, value is set, state changes to
|
||||
// COMPLETED, and done is closed, unblocking waiters.
|
||||
//
|
||||
// Alternatively, as Get calls are cancelled, they decrement
|
||||
// waiters. If it drops to zero, the inner context is
|
||||
// cancelled, computation is abandoned, and state resets to
|
||||
// IDLE to start the process over again.
|
||||
state state
|
||||
// done is set in running state, and closed when exiting it.
|
||||
done chan struct{}
|
||||
// cancel is set in running state. It cancels computation.
|
||||
cancel context.CancelFunc
|
||||
// waiters is the number of Gets outstanding.
|
||||
waiters uint
|
||||
// the function that will be used to populate the value
|
||||
function Function
|
||||
// value is set in completed state.
|
||||
value interface{}
|
||||
}
|
||||
|
||||
// NewPromise returns a promise for the future result of calling the
|
||||
// specified function.
|
||||
//
|
||||
// The debug string is used to classify promises in logs and metrics.
|
||||
// It should be drawn from a small set.
|
||||
func NewPromise(debug string, function Function) *Promise {
|
||||
if function == nil {
|
||||
panic("nil function")
|
||||
}
|
||||
return &Promise{
|
||||
debug: debug,
|
||||
function: function,
|
||||
}
|
||||
}
|
||||
|
||||
type state int
|
||||
|
||||
const (
|
||||
stateIdle = iota // newly constructed, or last waiter was cancelled
|
||||
stateRunning // start was called and not cancelled
|
||||
stateCompleted // function call ran to completion
|
||||
)
|
||||
|
||||
// Cached returns the value associated with a promise.
|
||||
//
|
||||
// It will never cause the value to be generated.
|
||||
// It will return the cached value, if present.
|
||||
func (p *Promise) Cached() interface{} {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.state == stateCompleted {
|
||||
return p.value
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get returns the value associated with a promise.
|
||||
//
|
||||
// All calls to Promise.Get on a given promise return the
|
||||
// same result but the function is called (to completion) at most once.
|
||||
//
|
||||
// If the value is not yet ready, the underlying function will be invoked.
|
||||
//
|
||||
// If ctx is cancelled, Get returns (nil, Canceled).
|
||||
// If all concurrent calls to Get are cancelled, the context provided
|
||||
// to the function is cancelled. A later call to Get may attempt to
|
||||
// call the function again.
|
||||
func (p *Promise) Get(ctx context.Context, arg interface{}) (interface{}, error) {
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
p.mu.Lock()
|
||||
switch p.state {
|
||||
case stateIdle:
|
||||
return p.run(ctx, arg)
|
||||
case stateRunning:
|
||||
return p.wait(ctx)
|
||||
case stateCompleted:
|
||||
defer p.mu.Unlock()
|
||||
return p.value, nil
|
||||
default:
|
||||
panic("unknown state")
|
||||
}
|
||||
}
|
||||
|
||||
// run starts p.function and returns the result. p.mu must be locked.
|
||||
func (p *Promise) run(ctx context.Context, arg interface{}) (interface{}, error) {
|
||||
childCtx, cancel := context.WithCancel(xcontext.Detach(ctx))
|
||||
p.cancel = cancel
|
||||
p.state = stateRunning
|
||||
p.done = make(chan struct{})
|
||||
function := p.function // Read under the lock
|
||||
|
||||
// Make sure that the argument isn't destroyed while we're running in it.
|
||||
release := func() {}
|
||||
if rc, ok := arg.(RefCounted); ok {
|
||||
release = rc.Acquire()
|
||||
}
|
||||
|
||||
go func() {
|
||||
trace.WithRegion(childCtx, fmt.Sprintf("Promise.run %s", p.debug), func() {
|
||||
defer release()
|
||||
// Just in case the function does something expensive without checking
|
||||
// the context, double-check we're still alive.
|
||||
if childCtx.Err() != nil {
|
||||
return
|
||||
}
|
||||
v := function(childCtx, arg)
|
||||
if childCtx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
// It's theoretically possible that the promise has been cancelled out
|
||||
// of the run that started us, and then started running again since we
|
||||
// checked childCtx above. Even so, that should be harmless, since each
|
||||
// run should produce the same results.
|
||||
if p.state != stateRunning {
|
||||
return
|
||||
}
|
||||
|
||||
p.value = v
|
||||
p.function = nil // aid GC
|
||||
p.state = stateCompleted
|
||||
close(p.done)
|
||||
})
|
||||
}()
|
||||
|
||||
return p.wait(ctx)
|
||||
}
|
||||
|
||||
// wait waits for the value to be computed, or ctx to be cancelled. p.mu must be locked.
|
||||
func (p *Promise) wait(ctx context.Context) (interface{}, error) {
|
||||
p.waiters++
|
||||
done := p.done
|
||||
p.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.state == stateCompleted {
|
||||
return p.value, nil
|
||||
}
|
||||
return nil, nil
|
||||
case <-ctx.Done():
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.waiters--
|
||||
if p.waiters == 0 && p.state == stateRunning {
|
||||
p.cancel()
|
||||
close(p.done)
|
||||
p.state = stateIdle
|
||||
p.done = nil
|
||||
p.cancel = nil
|
||||
}
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// An EvictionPolicy controls the eviction behavior of keys in a Store when
|
||||
// they no longer have any references.
|
||||
type EvictionPolicy int
|
||||
|
||||
const (
|
||||
// ImmediatelyEvict evicts keys as soon as they no longer have references.
|
||||
ImmediatelyEvict EvictionPolicy = iota
|
||||
|
||||
// NeverEvict does not evict keys.
|
||||
NeverEvict
|
||||
)
|
||||
|
||||
// A Store maps arbitrary keys to reference-counted promises.
|
||||
//
|
||||
// The zero value is a valid Store, though a store may also be created via
|
||||
// NewStore if a custom EvictionPolicy is required.
|
||||
type Store struct {
|
||||
evictionPolicy EvictionPolicy
|
||||
|
||||
promisesMu sync.Mutex
|
||||
promises map[interface{}]*Promise
|
||||
}
|
||||
|
||||
// NewStore creates a new store with the given eviction policy.
|
||||
func NewStore(policy EvictionPolicy) *Store {
|
||||
return &Store{evictionPolicy: policy}
|
||||
}
|
||||
|
||||
// Promise returns a reference-counted promise for the future result of
|
||||
// calling the specified function.
|
||||
//
|
||||
// Calls to Promise with the same key return the same promise, incrementing its
|
||||
// reference count. The caller must call the returned function to decrement
|
||||
// the promise's reference count when it is no longer needed. The returned
|
||||
// function must not be called more than once.
|
||||
//
|
||||
// Once the last reference has been released, the promise is removed from the
|
||||
// store.
|
||||
func (store *Store) Promise(key interface{}, function Function) (*Promise, func()) {
|
||||
store.promisesMu.Lock()
|
||||
p, ok := store.promises[key]
|
||||
if !ok {
|
||||
p = NewPromise(reflect.TypeOf(key).String(), function)
|
||||
if store.promises == nil {
|
||||
store.promises = map[interface{}]*Promise{}
|
||||
}
|
||||
store.promises[key] = p
|
||||
}
|
||||
p.refcount++
|
||||
store.promisesMu.Unlock()
|
||||
|
||||
var released int32
|
||||
release := func() {
|
||||
if !atomic.CompareAndSwapInt32(&released, 0, 1) {
|
||||
panic("release called more than once")
|
||||
}
|
||||
store.promisesMu.Lock()
|
||||
|
||||
p.refcount--
|
||||
if p.refcount == 0 && store.evictionPolicy != NeverEvict {
|
||||
// Inv: if p.refcount > 0, then store.promises[key] == p.
|
||||
delete(store.promises, key)
|
||||
}
|
||||
store.promisesMu.Unlock()
|
||||
}
|
||||
|
||||
return p, release
|
||||
}
|
||||
|
||||
// Stats returns the number of each type of key in the store.
|
||||
func (s *Store) Stats() map[reflect.Type]int {
|
||||
result := map[reflect.Type]int{}
|
||||
|
||||
s.promisesMu.Lock()
|
||||
defer s.promisesMu.Unlock()
|
||||
|
||||
for k := range s.promises {
|
||||
result[reflect.TypeOf(k)]++
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// DebugOnlyIterate iterates through the store and, for each completed
|
||||
// promise, calls f(k, v) for the map key k and function result v. It
|
||||
// should only be used for debugging purposes.
|
||||
func (s *Store) DebugOnlyIterate(f func(k, v interface{})) {
|
||||
s.promisesMu.Lock()
|
||||
defer s.promisesMu.Unlock()
|
||||
|
||||
for k, p := range s.promises {
|
||||
if v := p.Cached(); v != nil {
|
||||
f(k, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
// Copyright 2019 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package memoize_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/tools/internal/memoize"
|
||||
)
|
||||
|
||||
func TestGet(t *testing.T) {
|
||||
var store memoize.Store
|
||||
|
||||
evaled := 0
|
||||
|
||||
h, release := store.Promise("key", func(context.Context, interface{}) interface{} {
|
||||
evaled++
|
||||
return "res"
|
||||
})
|
||||
defer release()
|
||||
expectGet(t, h, "res")
|
||||
expectGet(t, h, "res")
|
||||
if evaled != 1 {
|
||||
t.Errorf("got %v calls to function, wanted 1", evaled)
|
||||
}
|
||||
}
|
||||
|
||||
func expectGet(t *testing.T, h *memoize.Promise, wantV interface{}) {
|
||||
t.Helper()
|
||||
gotV, gotErr := h.Get(context.Background(), nil)
|
||||
if gotV != wantV || gotErr != nil {
|
||||
t.Fatalf("Get() = %v, %v, wanted %v, nil", gotV, gotErr, wantV)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewPromise(t *testing.T) {
|
||||
calls := 0
|
||||
f := func(context.Context, interface{}) interface{} {
|
||||
calls++
|
||||
return calls
|
||||
}
|
||||
|
||||
// All calls to Get on the same promise return the same result.
|
||||
p1 := memoize.NewPromise("debug", f)
|
||||
expectGet(t, p1, 1)
|
||||
expectGet(t, p1, 1)
|
||||
|
||||
// A new promise calls the function again.
|
||||
p2 := memoize.NewPromise("debug", f)
|
||||
expectGet(t, p2, 2)
|
||||
expectGet(t, p2, 2)
|
||||
|
||||
// The original promise is unchanged.
|
||||
expectGet(t, p1, 1)
|
||||
}
|
||||
|
||||
func TestStoredPromiseRefCounting(t *testing.T) {
|
||||
var store memoize.Store
|
||||
v1 := false
|
||||
v2 := false
|
||||
p1, release1 := store.Promise("key1", func(context.Context, interface{}) interface{} {
|
||||
return &v1
|
||||
})
|
||||
p2, release2 := store.Promise("key2", func(context.Context, interface{}) interface{} {
|
||||
return &v2
|
||||
})
|
||||
expectGet(t, p1, &v1)
|
||||
expectGet(t, p2, &v2)
|
||||
|
||||
expectGet(t, p1, &v1)
|
||||
expectGet(t, p2, &v2)
|
||||
|
||||
p2Copy, release2Copy := store.Promise("key2", func(context.Context, interface{}) interface{} {
|
||||
return &v1
|
||||
})
|
||||
if p2 != p2Copy {
|
||||
t.Error("Promise returned a new value while old is not destroyed yet")
|
||||
}
|
||||
expectGet(t, p2Copy, &v2)
|
||||
|
||||
release2()
|
||||
if got, want := v2, false; got != want {
|
||||
t.Errorf("after destroying first v2 ref, got %v, want %v", got, want)
|
||||
}
|
||||
release2Copy()
|
||||
if got, want := v1, false; got != want {
|
||||
t.Errorf("after destroying v2, got %v, want %v", got, want)
|
||||
}
|
||||
release1()
|
||||
|
||||
p2Copy, release2Copy = store.Promise("key2", func(context.Context, interface{}) interface{} {
|
||||
return &v2
|
||||
})
|
||||
if p2 == p2Copy {
|
||||
t.Error("Promise returned previously destroyed value")
|
||||
}
|
||||
release2Copy()
|
||||
}
|
||||
|
||||
func TestPromiseDestroyedWhileRunning(t *testing.T) {
|
||||
// Test that calls to Promise.Get return even if the promise is destroyed while running.
|
||||
|
||||
var store memoize.Store
|
||||
c := make(chan int)
|
||||
|
||||
var v int
|
||||
h, release := store.Promise("key", func(ctx context.Context, _ interface{}) interface{} {
|
||||
<-c
|
||||
<-c
|
||||
if err := ctx.Err(); err != nil {
|
||||
t.Errorf("ctx.Err() = %v, want nil", err)
|
||||
}
|
||||
return &v
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) // arbitrary timeout; may be removed if it causes flakes
|
||||
defer cancel()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
var got interface{}
|
||||
var err error
|
||||
go func() {
|
||||
got, err = h.Get(ctx, nil)
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
c <- 0 // send once to enter the promise function
|
||||
release() // release before the promise function returns
|
||||
c <- 0 // let the promise function proceed
|
||||
|
||||
wg.Wait()
|
||||
|
||||
if err != nil {
|
||||
t.Errorf("Get() failed: %v", err)
|
||||
}
|
||||
if got != &v {
|
||||
t.Errorf("Get() = %v, want %v", got, v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoubleReleasePanics(t *testing.T) {
|
||||
var store memoize.Store
|
||||
_, release := store.Promise("key", func(ctx context.Context, _ interface{}) interface{} { return 0 })
|
||||
|
||||
panicked := false
|
||||
|
||||
func() {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
panicked = true
|
||||
}
|
||||
}()
|
||||
release()
|
||||
release()
|
||||
}()
|
||||
|
||||
if !panicked {
|
||||
t.Errorf("calling release() twice did not panic")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user