internal/memoize: don't destroy reference counted handles

Unlike generational handles, when reference counted handles are evicted
from the Store we don't know that they are also no longer in use by
active goroutines. Destroying them causes goroutine leaks.

Also fix a data race because Handle.mu was not acquired in the release
func returned by GetHandle.

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

internal/lsp/cache/parse.go

@@ -61,7 +61,7 @@ func (s *snapshot) parseGoHandle(ctx context.Context, fh source.FileHandle, mode
 	parseHandle, release := s.generation.GetHandle(key, func(ctx context.Context, arg memoize.Arg) interface{} {
 		snapshot := arg.(*snapshot)
 		return parseGo(ctx, snapshot.FileSet(), fh, mode)
-	}, nil)
+	})
 
 	pgh := &parseGoHandle{
 		handle: parseHandle,

internal/memoize/memoize.go

@@ -83,18 +83,18 @@ func (g *Generation) Destroy(destroyedBy string) {
 
 	g.store.mu.Lock()
 	defer g.store.mu.Unlock()
-	for _, e := range g.store.handles {
-		if !e.trackGenerations {
+	for _, h := range g.store.handles {
+		if !h.trackGenerations {
 			continue
 		}
-		e.mu.Lock()
-		if _, ok := e.generations[g]; ok {
-			delete(e.generations, g) // delete even if it's dead, in case of dangling references to the entry.
-			if len(e.generations) == 0 {
-				e.destroy(g.store)
+		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.
+			if len(h.generations) == 0 {
+				h.destroy(g.store)
 			}
 		}
-		e.mu.Unlock()
+		h.mu.Unlock()
 	}
 	delete(g.store.generations, g)
 }
@@ -120,6 +120,12 @@ type Function func(ctx context.Context, arg Arg) interface{}
 
 type state int
 
+// TODO(rfindley): remove stateDestroyed; Handles should not need to know
+// whether or not they have been destroyed.
+//
+// TODO(rfindley): also consider removing stateIdle. Why create a handle if you
+// aren't certain you're going to need its result? And if you know you need its
+// result, why wait to begin computing it?
 const (
 	stateIdle = iota
 	stateRunning
@@ -139,6 +145,12 @@ const (
 // 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.
+//
+// Handles may be tracked by generations, or directly reference counted, as
+// determined by the trackGenerations field. See the field comments for more
+// information about the differences between these two forms.
+//
+// TODO(rfindley): eliminate generational handles.
 type Handle struct {
 	key interface{}
 	mu  sync.Mutex
@@ -159,6 +171,11 @@ type Handle struct {
 	value interface{}
 	// cleanup, if non-nil, is used to perform any necessary clean-up on values
 	// produced by function.
+	//
+	// cleanup is never set for reference counted handles.
+	//
+	// TODO(rfindley): remove this field once workspace folders no longer need to
+	// be tracked.
 	cleanup func(interface{})
 
 	// If trackGenerations is set, this handle tracks generations in which it
@@ -190,19 +207,27 @@ func (g *Generation) Bind(key interface{}, function Function, cleanup func(inter
 //
 // As in opposite to Bind it returns a release callback which has to be called
 // once this reference to handle is not needed anymore.
-func (g *Generation) GetHandle(key interface{}, function Function, cleanup func(interface{})) (*Handle, func()) {
-	handle := g.getHandle(key, function, cleanup, false)
+func (g *Generation) GetHandle(key interface{}, function Function) (*Handle, func()) {
+	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()
 
-		handle.refCounter--
-		if handle.refCounter == 0 {
-			handle.destroy(store)
+		h.mu.Lock()
+		defer h.mu.Unlock()
+
+		h.refCounter--
+		if h.refCounter == 0 {
+			// Don't call h.destroy: for reference counted handles we can't know when
+			// they are no longer reachable from runnable goroutines. For example,
+			// gopls could have a current operation that is using a packageHandle.
+			// Destroying the handle here would cause that operation to hang.
+			delete(store.handles, h.key)
 		}
 	}
-	return handle, release
+	return h, release
 }
 
 func (g *Generation) getHandle(key interface{}, function Function, cleanup func(interface{}), trackGenerations bool) *Handle {
@@ -252,13 +277,13 @@ func (s *Store) DebugOnlyIterate(f func(k, v interface{})) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	for k, e := range s.handles {
+	for k, h := range s.handles {
 		var v interface{}
-		e.mu.Lock()
-		if e.state == stateCompleted {
-			v = e.value
+		h.mu.Lock()
+		if h.state == stateCompleted {
+			v = h.value
 		}
-		e.mu.Unlock()
+		h.mu.Unlock()
 		if v == nil {
 			continue
 		}
@@ -278,6 +303,7 @@ func (g *Generation) Inherit(h *Handle) {
 	h.incrementRef(g)
 }
 
+// destroy marks h as destroyed. h.mu and store.mu must be held.
 func (h *Handle) destroy(store *Store) {
 	h.state = stateDestroyed
 	if h.cleanup != nil && h.value != nil {
@@ -409,6 +435,12 @@ func (h *Handle) run(ctx context.Context, g *Generation, arg Arg) (interface{},
 				}
 				return
 			}
+
+			if h.cleanup != nil && h.value != nil {
+				// Clean up before overwriting an existing value.
+				h.cleanup(h.value)
+			}
+
 			// At this point v will be cleaned up whenever h is destroyed.
 			h.value = v
 			h.function = nil

internal/memoize/memoize_test.go

@@ -7,7 +7,9 @@ package memoize_test
 import (
 	"context"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 
 	"golang.org/x/tools/internal/memoize"
 )
@@ -112,15 +114,12 @@ func TestHandleRefCounting(t *testing.T) {
 	g1 := s.Generation("g1")
 	v1 := false
 	v2 := false
-	cleanup := func(v interface{}) {
-		*(v.(*bool)) = true
-	}
 	h1, release1 := g1.GetHandle("key1", func(context.Context, memoize.Arg) interface{} {
 		return &v1
-	}, nil)
+	})
 	h2, release2 := g1.GetHandle("key2", func(context.Context, memoize.Arg) interface{} {
 		return &v2
-	}, cleanup)
+	})
 	expectGet(t, h1, g1, &v1)
 	expectGet(t, h2, g1, &v2)
 
@@ -131,7 +130,7 @@ func TestHandleRefCounting(t *testing.T) {
 
 	h2Copy, release2Copy := g2.GetHandle("key2", func(context.Context, memoize.Arg) interface{} {
 		return &v1
-	}, nil)
+	})
 	if h2 != h2Copy {
 		t.Error("NewHandle returned a new value while old is not destroyed yet")
 	}
@@ -140,24 +139,65 @@ func TestHandleRefCounting(t *testing.T) {
 
 	release2()
 	if got, want := v2, false; got != want {
-		t.Error("after destroying first v2 ref, v2 is cleaned up")
+		t.Errorf("after destroying first v2 ref, got %v, want %v", got, want)
 	}
 	release2Copy()
-	if got, want := v2, true; got != want {
-		t.Error("after destroying second v2 ref, v2 is not cleaned up")
-	}
 	if got, want := v1, false; got != want {
-		t.Error("after destroying v2, v1 is cleaned up")
+		t.Errorf("after destroying v2, got %v, want %v", got, want)
 	}
 	release1()
 
 	g3 := s.Generation("g3")
 	h2Copy, release2Copy = g3.GetHandle("key2", func(context.Context, memoize.Arg) interface{} {
 		return &v2
-	}, cleanup)
+	})
 	if h2 == h2Copy {
 		t.Error("NewHandle returned previously destroyed value")
 	}
 	release2Copy()
 	g3.Destroy("by test")
 }
+
+func TestHandleDestroyedWhileRunning(t *testing.T) {
+	// Test that calls to Handle.Get return even if the handle is destroyed while
+	// running.
+
+	s := &memoize.Store{}
+	g := s.Generation("g")
+	c := make(chan int)
+
+	var v int
+	h, release := g.GetHandle("key", func(ctx context.Context, _ memoize.Arg) 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, g, nil)
+		wg.Done()
+	}()
+
+	c <- 0    // send once to enter the handle function
+	release() // release before the handle function returns
+	c <- 0    // let the handle function proceed
+
+	wg.Wait()
+
+	if err != nil {
+		t.Errorf("Get() failed: %v", err)
+	}
+	if got != &v {
+		t.Errorf("Get() = %v, want %v", got, v)
+	}
+}