gopls/internal/server: avoid duplicate diagnoses and loads

With gopls@v0.15.0, zero config gopls made it much more likely that
sessions would have multiple Views. Additionally, with improved build
tag support, it made it more likely that these Views would share files.
As a result, we encountered (and fixed) this latent bug:

1. User changes file x.go, invalidating view A and B. A and B are
   scheduled for diagnosis.
2. User changes file y.go, invalidating only view B. Step (1) is
   cancelled and view B is scheduled for diagnosis.
3. View A never gets rediagnosed.

The fix was naive: just mark view A and B as dirty, and schedule a
goroutine to diagnose all dirty views after each step. As before, step
(2) would cancel the context from step (1).

But there's a problem: diagnoses were happening on the *Snapshot*
context, not the operation context. Therefore, if the goroutines of step
(1) and (2) both find the same snapshots, the diagnostics of step (1)
would *not* be cancelled, and would be performed in addition to the
diagnostics of (2). In other words, following a sequence of
invalidations, we could theoretically be collecting diagnostics N times
rather than 1 time.

In practice, this is not so much of a problem for smaller repositories.
Most of the time, changes are arriving at the speed of keystrokes, and
diagnostics are being scheduled faster than we can type. However, on
slower machines, or when there is more overall work being scheduled, or
when changes arrive simultaneously (such as with a 'save all' command or
branch switch), it is quite possible in practice to cause gopls to do
more work than necessary, including redundant loads. I'm not sure if
this is what conspires to cause the regressions described in
golang/go#66647, but it certainly is a real regression.

Fix this by threading the correct context into diagnoseSnapshot.
Additionally, add earlier context cancellation in a few cases where
redundant work was being performed despite a context cancellation.

For golang/go#66647

Change-Id: I67da1c186848286ca7b6221330a655d23820fd5d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/577695
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>

gopls/internal/cache/load.go

@@ -43,6 +43,11 @@ var errNoPackages = errors.New("no packages returned")
 //
 // If scopes contains a file scope there must be exactly one scope.
 func (s *Snapshot) load(ctx context.Context, allowNetwork bool, scopes ...loadScope) (err error) {
+	if ctx.Err() != nil {
+		// Check context cancellation before incrementing id below: a load on a
+		// cancelled context should be a no-op.
+		return ctx.Err()
+	}
 	id := atomic.AddUint64(&loadID, 1)
 	eventName := fmt.Sprintf("go/packages.Load #%d", id) // unique name for logging
 

gopls/internal/cache/session.go

@@ -369,11 +369,11 @@ func (s *Session) SnapshotOf(ctx context.Context, uri protocol.DocumentURI) (*Sn
 		if err != nil {
 			continue // view was shut down
 		}
-		_ = snapshot.awaitLoaded(ctx) // ignore error
-		g := snapshot.MetadataGraph()
 		// We don't check the error from awaitLoaded, because a load failure (that
 		// doesn't result from context cancelation) should not prevent us from
 		// continuing to search for the best view.
+		_ = snapshot.awaitLoaded(ctx)
+		g := snapshot.MetadataGraph()
 		if ctx.Err() != nil {
 			release()
 			return nil, nil, ctx.Err()
@@ -1121,6 +1121,11 @@ func (s *Session) FileWatchingGlobPatterns(ctx context.Context) map[protocol.Rel
 //
 // The caller must not mutate the result.
 func (s *Session) OrphanedFileDiagnostics(ctx context.Context) (map[protocol.DocumentURI][]*Diagnostic, error) {
+	if err := ctx.Err(); err != nil {
+		// Avoid collecting diagnostics if the context is cancelled.
+		// (Previously, it was possible to get all the way to packages.Load on a cancelled context)
+		return nil, err
+	}
 	// Note: diagnostics holds a slice for consistency with other diagnostic
 	// funcs.
 	diagnostics := make(map[protocol.DocumentURI][]*Diagnostic)

gopls/internal/cache/snapshot.go

@@ -1381,6 +1381,10 @@ func (s *Snapshot) AwaitInitialized(ctx context.Context) {
 
 // reloadWorkspace reloads the metadata for all invalidated workspace packages.
 func (s *Snapshot) reloadWorkspace(ctx context.Context) {
+	if ctx.Err() != nil {
+		return
+	}
+
 	var scopes []loadScope
 	var seen map[PackagePath]bool
 	s.mu.Lock()

gopls/internal/server/command.go

@@ -281,7 +281,9 @@ func (c *commandHandler) modifyState(ctx context.Context, source ModificationSou
 	}
 	wg.Add(1)
 	go func() {
-		c.s.diagnoseSnapshot(snapshot, nil, 0)
+		// Diagnosing with the background context ensures new snapshots are fully
+		// diagnosed.
+		c.s.diagnoseSnapshot(snapshot.BackgroundContext(), snapshot, nil, 0)
 		release()
 		wg.Done()
 	}()
@@ -1076,7 +1078,10 @@ func (c *commandHandler) RunGovulncheck(ctx context.Context, args command.Vulnch
 			return err
 		}
 		defer release()
-		c.s.diagnoseSnapshot(snapshot, nil, 0)
+
+		// Diagnosing with the background context ensures new snapshots are fully
+		// diagnosed.
+		c.s.diagnoseSnapshot(snapshot.BackgroundContext(), snapshot, nil, 0)
 
 		affecting := make(map[string]bool, len(result.Entries))
 		for _, finding := range result.Findings {
@@ -1408,7 +1413,11 @@ func (c *commandHandler) DiagnoseFiles(ctx context.Context, args command.Diagnos
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
-				c.s.diagnoseSnapshot(snapshot, nil, 0)
+
+				// Use the operation context for diagnosis, rather than
+				// snapshot.BackgroundContext, because this operation does not create
+				// new snapshots (so they should also be diagnosed by other means).
+				c.s.diagnoseSnapshot(ctx, snapshot, nil, 0)
 			}()
 		}
 		wg.Wait()

gopls/internal/server/diagnostics.go

@@ -135,16 +135,16 @@ func (s *server) diagnoseChangedViews(ctx context.Context, modID uint64, lastCha
 		go func(snapshot *cache.Snapshot, uris []protocol.DocumentURI) {
 			defer release()
 			defer wg.Done()
-			s.diagnoseSnapshot(snapshot, uris, snapshot.Options().DiagnosticsDelay)
+			s.diagnoseSnapshot(ctx, snapshot, uris, snapshot.Options().DiagnosticsDelay)
 			s.modificationMu.Lock()
 
-			// Only remove v from s.viewsToDiagnose if the snapshot is not cancelled.
+			// Only remove v from s.viewsToDiagnose if the context is not cancelled.
 			// This ensures that the snapshot was not cloned before its state was
 			// fully evaluated, and therefore avoids missing a change that was
 			// irrelevant to an incomplete snapshot.
 			//
 			// See the documentation for s.viewsToDiagnose for details.
-			if snapshot.BackgroundContext().Err() == nil && s.viewsToDiagnose[v] <= modID {
+			if ctx.Err() == nil && s.viewsToDiagnose[v] <= modID {
 				delete(s.viewsToDiagnose, v)
 			}
 			s.modificationMu.Unlock()
@@ -173,8 +173,14 @@ func (s *server) diagnoseChangedViews(ctx context.Context, modID uint64, lastCha
 // If changedURIs is non-empty, it is a set of recently changed files that
 // should be diagnosed immediately, and onDisk reports whether these file
 // changes came from a change to on-disk files.
-func (s *server) diagnoseSnapshot(snapshot *cache.Snapshot, changedURIs []protocol.DocumentURI, delay time.Duration) {
-	ctx := snapshot.BackgroundContext()
+//
+// If the provided context is cancelled, diagnostics may be partially
+// published. Therefore, the provided context should only be cancelled if there
+// will be a subsequent operation to make diagnostics consistent. In general,
+// if an operation creates a new snapshot, it is responsible for ensuring that
+// snapshot (or a subsequent snapshot in the same View) is eventually
+// diagnosed.
+func (s *server) diagnoseSnapshot(ctx context.Context, snapshot *cache.Snapshot, changedURIs []protocol.DocumentURI, delay time.Duration) {
 	ctx, done := event.Start(ctx, "Server.diagnoseSnapshot", snapshot.Labels()...)
 	defer done()
 

gopls/internal/server/general.go

@@ -332,7 +332,7 @@ func (s *server) addFolders(ctx context.Context, folders []protocol.WorkspaceFol
 		// Diagnose the newly created view asynchronously.
 		ndiagnose.Add(1)
 		go func() {
-			s.diagnoseSnapshot(snapshot, nil, 0)
+			s.diagnoseSnapshot(snapshot.BackgroundContext(), snapshot, nil, 0)
 			<-initialized
 			release()
 			ndiagnose.Done()