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>
diff --git a/internal/lsp/cache/imports.go b/internal/lsp/cache/imports.go
index ed9919f..b8e0146 100644
--- a/internal/lsp/cache/imports.go
+++ b/internal/lsp/cache/imports.go
@@ -138,7 +138,7 @@
 
 	// 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(),
 	})
diff --git a/internal/lsp/cache/session.go b/internal/lsp/cache/session.go
index 48acf46..a65b8fe 100644
--- a/internal/lsp/cache/session.go
+++ b/internal/lsp/cache/session.go
@@ -254,7 +254,7 @@
 	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 @@
 			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) 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 @@
 	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 @@
 	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
 		}
diff --git a/internal/lsp/cache/view.go b/internal/lsp/cache/view.go
index 91c9108..fcff02a 100644
--- a/internal/lsp/cache/view.go
+++ b/internal/lsp/cache/view.go
@@ -74,7 +74,7 @@
 	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 @@
 	}
 	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 (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 @@
 
 // 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 @@
 	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 @@
 	}
 	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 @@
 // 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 {
diff --git a/internal/memoize/memoize.go b/internal/memoize/memoize.go
index 6127fc3..0037342 100644
--- a/internal/memoize/memoize.go
+++ b/internal/memoize/memoize.go
@@ -102,11 +102,8 @@
 
 // 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 @@
 // 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 @@
 	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