internal/memoize: do not iterate all handles on generation destruction

This allows reducing critical section of `g.store.mu` as the vast
majority of entries do not rely on generation-based GC.

Change-Id: I985af0b38504ddedb22649290deac91797577b75
Reviewed-on: https://go-review.googlesource.com/c/tools/+/413656
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>

internal/memoize/memoize.go

@@ -33,6 +33,8 @@ var (
 type Store struct {
 	handlesMu sync.Mutex // lock ordering: Store.handlesMu before Handle.mu
 	handles   map[interface{}]*Handle
+	// handles which are bound to generations for GC purposes.
+	boundHandles map[*Handle]struct{}
 }
 
 // Generation creates a new Generation associated with s. Destroy must be
@@ -71,10 +73,7 @@ func (g *Generation) Destroy(destroyedBy string) {
 
 	g.store.handlesMu.Lock()
 	defer g.store.handlesMu.Unlock()
-	for _, h := range g.store.handles {
-		if !h.trackGenerations {
-			continue
-		}
+	for h := range g.store.boundHandles {
 		h.mu.Lock()
 		if _, ok := h.generations[g]; ok {
 			delete(h.generations, g) // delete even if it's dead, in case of dangling references to the entry.
@@ -237,7 +236,11 @@ func (g *Generation) getHandle(key interface{}, function Function, cleanup func(
 			trackGenerations: trackGenerations,
 		}
 		if trackGenerations {
+			if g.store.boundHandles == nil {
+				g.store.boundHandles = map[*Handle]struct{}{}
+			}
 			h.generations = make(map[*Generation]struct{}, 1)
+			g.store.boundHandles[h] = struct{}{}
 		}
 
 		if g.store.handles == nil {
@@ -302,6 +305,9 @@ func (h *Handle) destroy(store *Store) {
 		h.cleanup(h.value)
 	}
 	delete(store.handles, h.key)
+	if h.trackGenerations {
+		delete(store.boundHandles, h)
+	}
 }
 
 func (h *Handle) incrementRef(g *Generation) {