diff --git a/internal/memoize/memoize.go b/internal/memoize/memoize.go index 6477abbb81..a758deeb7f 100644 --- a/internal/memoize/memoize.go +++ b/internal/memoize/memoize.go @@ -31,27 +31,15 @@ var ( // Store binds keys to functions, returning handles that can be used to access // the functions results. type Store struct { - mu sync.Mutex - // handles is the set of values stored. - handles map[interface{}]*Handle - - // generations is the set of generations live in this store. - generations map[*Generation]struct{} + handlesMu sync.Mutex // lock ordering: Store.handlesMu before Handle.mu + handles map[interface{}]*Handle } // Generation creates a new Generation associated with s. Destroy must be // called on the returned Generation once it is no longer in use. name is // for debugging purposes only. func (s *Store) Generation(name string) *Generation { - s.mu.Lock() - defer s.mu.Unlock() - if s.handles == nil { - s.handles = map[interface{}]*Handle{} - s.generations = map[*Generation]struct{}{} - } - g := &Generation{store: s, name: name} - s.generations[g] = struct{}{} - return g + return &Generation{store: s, name: name} } // A Generation is a logical point in time of the cache life-cycle. Cache @@ -81,8 +69,8 @@ func (g *Generation) Destroy(destroyedBy string) { panic("Destroy on generation " + g.name + " already destroyed by " + prevDestroyedBy) } - g.store.mu.Lock() - defer g.store.mu.Unlock() + g.store.handlesMu.Lock() + defer g.store.handlesMu.Unlock() for _, h := range g.store.handles { if !h.trackGenerations { continue @@ -96,7 +84,6 @@ func (g *Generation) Destroy(destroyedBy string) { } h.mu.Unlock() } - delete(g.store.generations, g) } // Acquire creates a new reference to g, and returns a func to release that @@ -153,7 +140,7 @@ const ( // TODO(rfindley): eliminate generational handles. type Handle struct { key interface{} - mu sync.Mutex + mu sync.Mutex // lock ordering: Store.handlesMu before Handle.mu // generations is the set of generations in which this handle is valid. generations map[*Generation]struct{} @@ -211,9 +198,9 @@ func (g *Generation) GetHandle(key interface{}, function Function) (*Handle, fun h := g.getHandle(key, function, nil, false) store := g.store release := func() { - // Acquire store.mu before mutating refCounter - store.mu.Lock() - defer store.mu.Unlock() + // Acquire store.handlesMu before mutating refCounter + store.handlesMu.Lock() + defer store.handlesMu.Unlock() h.mu.Lock() defer h.mu.Unlock() @@ -239,8 +226,8 @@ func (g *Generation) getHandle(key interface{}, function Function, cleanup func( if atomic.LoadUint32(&g.destroyed) != 0 { panic("operation on generation " + g.name + " destroyed by " + g.destroyedBy) } - g.store.mu.Lock() - defer g.store.mu.Unlock() + g.store.handlesMu.Lock() + defer g.store.handlesMu.Unlock() h, ok := g.store.handles[key] if !ok { h = &Handle{ @@ -252,6 +239,10 @@ func (g *Generation) getHandle(key interface{}, function Function, cleanup func( if trackGenerations { h.generations = make(map[*Generation]struct{}, 1) } + + if g.store.handles == nil { + g.store.handles = map[interface{}]*Handle{} + } g.store.handles[key] = h } @@ -261,10 +252,11 @@ func (g *Generation) getHandle(key interface{}, function Function, cleanup func( // Stats returns the number of each type of value in the store. func (s *Store) Stats() map[reflect.Type]int { - s.mu.Lock() - defer s.mu.Unlock() - result := map[reflect.Type]int{} + + s.handlesMu.Lock() + defer s.handlesMu.Unlock() + for k := range s.handles { result[reflect.TypeOf(k)]++ } @@ -274,8 +266,8 @@ func (s *Store) Stats() map[reflect.Type]int { // DebugOnlyIterate iterates through all live cache entries and calls f on them. // It should only be used for debugging purposes. func (s *Store) DebugOnlyIterate(f func(k, v interface{})) { - s.mu.Lock() - defer s.mu.Unlock() + s.handlesMu.Lock() + defer s.handlesMu.Unlock() for k, h := range s.handles { var v interface{} @@ -303,7 +295,7 @@ func (g *Generation) Inherit(h *Handle) { h.incrementRef(g) } -// destroy marks h as destroyed. h.mu and store.mu must be held. +// destroy marks h as destroyed. h.mu and store.handlesMu must be held. func (h *Handle) destroy(store *Store) { h.state = stateDestroyed if h.cleanup != nil && h.value != nil {