internal/memoize: do not allow (*Generation).Acquire to fail

The Acquire method is nearly instantaneous; it only potentially blocks
on a small, constant sequence of cache misses, so there is no need to
avoid blocking in it when a Context is cancelled.

An early return when the passed-in Context is canceled was added in CL
242838 to avoid incrementing the Generation's WaitGroup after its
destruction has begun; however, that early return also bypasses the
WaitGroup accounting that blocks Destroy while the generation is still
in use. Instead, we need the invariant that Acquire is not called in
the first place after Destroy, which we can ensure by nilling out
the View's snapshot when we begin destroying it.

I was not able to reproduce golang/go#48774 locally, but I believe
that this CL will fix it. (It may, however, expose other races or
deadlocks that may have been masked by the early return, which we can
then fix separately.)

Fixes golang/go#48774

Change-Id: Iac36fceb06485f849da5ba0250b44b55f937c44b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/367675
Trust: Bryan C. Mills <bcmills@google.com>
Run-TryBot: Bryan C. Mills <bcmills@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>
TryBot-Result: Go Bot <gobot@golang.org>

internal/lsp/cache/imports.go

@@ -138,7 +138,7 @@ func (s *importsState) populateProcessEnv(ctx context.Context, snapshot *snapsho
 
 	// Take an extra reference to the snapshot so that its workspace directory
 	// (if any) isn't destroyed while we're using it.
-	release := snapshot.generation.Acquire(ctx)
+	release := snapshot.generation.Acquire()
 	_, inv, cleanupInvocation, err := snapshot.goCommandInvocation(ctx, source.LoadWorkspace, &gocommand.Invocation{
 		WorkingDir: snapshot.view.rootURI.Filename(),
 	})

internal/lsp/cache/session.go

@@ -254,7 +254,7 @@ func (s *Session) createView(ctx context.Context, name string, folder, tempWorks
 	initCtx, initCancel := context.WithCancel(xcontext.Detach(ctx))
 	v.initCancelFirstAttempt = initCancel
 	snapshot := v.snapshot
-	release := snapshot.generation.Acquire(initCtx)
+	release := snapshot.generation.Acquire()
 	go func() {
 		defer release()
 		snapshot.initialize(initCtx, true)
@@ -264,7 +264,7 @@ func (s *Session) createView(ctx context.Context, name string, folder, tempWorks
 			event.Error(ctx, "copying workspace dir", err)
 		}
 	}()
-	return v, snapshot, snapshot.generation.Acquire(ctx), nil
+	return v, snapshot, snapshot.generation.Acquire(), nil
 }
 
 // View returns the view by name.
@@ -367,16 +367,24 @@ func (s *Session) removeView(ctx context.Context, view *View) error {
 func (s *Session) updateView(ctx context.Context, view *View, options *source.Options) (*View, error) {
 	s.viewMu.Lock()
 	defer s.viewMu.Unlock()
+
+	// Preserve the snapshot ID if we are recreating the view.
+	view.snapshotMu.Lock()
+	if view.snapshot == nil {
+		view.snapshotMu.Unlock()
+		panic("updateView called after View was already shut down")
+	}
+	snapshotID := view.snapshot.id
+	view.snapshotMu.Unlock()
+
 	i, err := s.dropView(ctx, view)
 	if err != nil {
 		return nil, err
 	}
-	// Preserve the snapshot ID if we are recreating the view.
-	view.snapshotMu.Lock()
-	snapshotID := view.snapshot.id
-	view.snapshotMu.Unlock()
+
 	v, _, release, err := s.createView(ctx, view.name, view.folder, view.tempWorkspace, options, snapshotID)
 	release()
+
 	if err != nil {
 		// we have dropped the old view, but could not create the new one
 		// this should not happen and is very bad, but we still need to clean
@@ -530,7 +538,7 @@ func (s *Session) ExpandModificationsToDirectories(ctx context.Context, changes
 	defer s.viewMu.RUnlock()
 	var snapshots []*snapshot
 	for _, v := range s.views {
-		snapshot, release := v.getSnapshot(ctx)
+		snapshot, release := v.getSnapshot()
 		defer release()
 		snapshots = append(snapshots, snapshot)
 	}
@@ -720,7 +728,7 @@ func (s *Session) FileWatchingGlobPatterns(ctx context.Context) map[string]struc
 	defer s.viewMu.RUnlock()
 	patterns := map[string]struct{}{}
 	for _, view := range s.views {
-		snapshot, release := view.getSnapshot(ctx)
+		snapshot, release := view.getSnapshot()
 		for k, v := range snapshot.fileWatchingGlobPatterns(ctx) {
 			patterns[k] = v
 		}

internal/lsp/cache/view.go

@@ -74,7 +74,7 @@ type View struct {
 	initCancelFirstAttempt context.CancelFunc
 
 	snapshotMu sync.Mutex
-	snapshot   *snapshot
+	snapshot   *snapshot // nil after shutdown has been called
 
 	// initialWorkspaceLoad is closed when the first workspace initialization has
 	// completed. If we failed to load, we only retry if the go.mod file changes,
@@ -505,7 +505,10 @@ func (v *View) shutdown(ctx context.Context) {
 	}
 	v.mu.Unlock()
 	v.snapshotMu.Lock()
-	go v.snapshot.generation.Destroy("View.shutdown")
+	if v.snapshot != nil {
+		go v.snapshot.generation.Destroy("View.shutdown")
+		v.snapshot = nil
+	}
 	v.snapshotMu.Unlock()
 	v.importsState.destroy()
 }
@@ -551,13 +554,16 @@ func checkIgnored(suffix string) bool {
 }
 
 func (v *View) Snapshot(ctx context.Context) (source.Snapshot, func()) {
-	return v.getSnapshot(ctx)
+	return v.getSnapshot()
 }
 
-func (v *View) getSnapshot(ctx context.Context) (*snapshot, func()) {
+func (v *View) getSnapshot() (*snapshot, func()) {
 	v.snapshotMu.Lock()
 	defer v.snapshotMu.Unlock()
-	return v.snapshot, v.snapshot.generation.Acquire(ctx)
+	if v.snapshot == nil {
+		panic("getSnapshot called after shutdown")
+	}
+	return v.snapshot, v.snapshot.generation.Acquire()
 }
 
 func (s *snapshot) initialize(ctx context.Context, firstAttempt bool) {
@@ -670,6 +676,9 @@ func (s *snapshot) loadWorkspace(ctx context.Context, firstAttempt bool) {
 
 // invalidateContent invalidates the content of a Go file,
 // including any position and type information that depends on it.
+//
+// invalidateContent returns a non-nil snapshot for the new content, along with
+// a callback which the caller must invoke to release that snapshot.
 func (v *View) invalidateContent(ctx context.Context, changes map[span.URI]*fileChange, forceReloadMetadata bool) (*snapshot, func()) {
 	// Detach the context so that content invalidation cannot be canceled.
 	ctx = xcontext.Detach(ctx)
@@ -678,6 +687,10 @@ func (v *View) invalidateContent(ctx context.Context, changes map[span.URI]*file
 	v.snapshotMu.Lock()
 	defer v.snapshotMu.Unlock()
 
+	if v.snapshot == nil {
+		panic("invalidateContent called after shutdown")
+	}
+
 	// Cancel all still-running previous requests, since they would be
 	// operating on stale data.
 	v.snapshot.cancel()
@@ -696,7 +709,7 @@ func (v *View) invalidateContent(ctx context.Context, changes map[span.URI]*file
 	}
 	go oldSnapshot.generation.Destroy("View.invalidateContent")
 
-	return v.snapshot, v.snapshot.generation.Acquire(ctx)
+	return v.snapshot, v.snapshot.generation.Acquire()
 }
 
 func (v *View) updateWorkspace(ctx context.Context) error {
@@ -714,7 +727,11 @@ func (v *View) updateWorkspace(ctx context.Context) error {
 // all changes to the workspace module, only that it is eventually consistent
 // with the workspace module of the latest snapshot.
 func (v *View) updateWorkspaceLocked(ctx context.Context) error {
-	release := v.snapshot.generation.Acquire(ctx)
+	if v.snapshot == nil {
+		return errors.New("view is shutting down")
+	}
+
+	release := v.snapshot.generation.Acquire()
 	defer release()
 	src, err := v.snapshot.getWorkspaceDir(ctx)
 	if err != nil {

internal/memoize/memoize.go

@@ -102,11 +102,8 @@ func (g *Generation) Destroy(destroyedBy string) {
 
 // Acquire creates a new reference to g, and returns a func to release that
 // reference.
-func (g *Generation) Acquire(ctx context.Context) func() {
+func (g *Generation) Acquire() func() {
 	destroyed := atomic.LoadUint32(&g.destroyed)
-	if ctx.Err() != nil {
-		return func() {}
-	}
 	if destroyed != 0 {
 		panic("acquire on generation " + g.name + " destroyed by " + g.destroyedBy)
 	}
@@ -274,7 +271,7 @@ func (h *Handle) Cached(g *Generation) interface{} {
 // If the value is not yet ready, the underlying function will be invoked.
 // If ctx is cancelled, Get returns nil.
 func (h *Handle) Get(ctx context.Context, g *Generation, arg Arg) (interface{}, error) {
-	release := g.Acquire(ctx)
+	release := g.Acquire()
 	defer release()
 
 	if ctx.Err() != nil {
@@ -319,7 +316,7 @@ func (h *Handle) run(ctx context.Context, g *Generation, arg Arg) (interface{},
 	function := h.function // Read under the lock
 
 	// Make sure that the generation isn't destroyed while we're running in it.
-	release := g.Acquire(ctx)
+	release := g.Acquire()
 	go func() {
 		defer release()
 		// Just in case the function does something expensive without checking