internal/lsp: diagnose the snapshot on every text synchronization event

This change moves to our ultimate approach of diagnostics the snapshot
on every file change, instead of carefully picking which files and
packages to diagnose. Analyses are shown for packages whose files are
open in the editor. Reverse dependencies are no longer needed for
source.Diagnostics because they will be invalidated when the snapshot is
cloned, so diagnosing the entire snapshot will bring them up to date.

This even works for go.mod files because all of workspace-level `go list`s
will be canceled as the user types, and then we trigger an uncancellable
go/packages.Load when the user saves. There is still room for improvement
here, but it will require much more careful invalidation of metadata for
go.mod files.

Change-Id: Id068505634b5e701c6f861a61b09a4c6704c565f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/214419
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
diff --git a/internal/lsp/cache/load.go b/internal/lsp/cache/load.go
index 19fa7c0..59eec94 100644
--- a/internal/lsp/cache/load.go
+++ b/internal/lsp/cache/load.go
@@ -85,8 +85,9 @@
 // shouldLoad reparses a file's package and import declarations to
 // determine if they have changed.
 func (c *cache) shouldLoad(ctx context.Context, s *snapshot, originalFH, currentFH source.FileHandle) bool {
+	// TODO(rstambler): go.mod files should be tracked in the snapshot.
 	if originalFH == nil {
-		return true
+		return currentFH.Identity().Kind == source.Go
 	}
 	// If the file hasn't changed, there's no need to reload.
 	if originalFH.Identity().String() == currentFH.Identity().String() {
diff --git a/internal/lsp/cache/snapshot.go b/internal/lsp/cache/snapshot.go
index 3785fcc..95b20bc 100644
--- a/internal/lsp/cache/snapshot.go
+++ b/internal/lsp/cache/snapshot.go
@@ -569,22 +569,34 @@
 
 // reloadWorkspace reloads the metadata for all invalidated workspace packages.
 func (s *snapshot) reloadWorkspace(ctx context.Context) error {
-	s.mu.Lock()
-	var scope []packagePath
-	for id, pkgPath := range s.workspacePackages {
-		if s.metadata[id] == nil {
-			scope = append(scope, pkgPath)
-		}
-	}
-	s.mu.Unlock()
-
-	if len(scope) == 0 {
+	scope := s.workspaceScope(ctx)
+	if scope == nil {
 		return nil
 	}
 	_, err := s.load(ctx, scope)
 	return err
 }
 
+func (s *snapshot) workspaceScope(ctx context.Context) interface{} {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	var pkgPaths []packagePath
+	for id, pkgPath := range s.workspacePackages {
+		if s.metadata[id] == nil {
+			pkgPaths = append(pkgPaths, pkgPath)
+		}
+	}
+	switch len(pkgPaths) {
+	case 0:
+		return nil
+	case len(s.workspacePackages):
+		return directoryURI(s.view.folder)
+	default:
+		return pkgPaths
+	}
+}
+
 func (s *snapshot) clone(ctx context.Context, withoutURI span.URI) *snapshot {
 	s.mu.Lock()
 	defer s.mu.Unlock()
diff --git a/internal/lsp/cmd/cmd_test.go b/internal/lsp/cmd/cmd_test.go
index 0e417c6..9b3f663 100644
--- a/internal/lsp/cmd/cmd_test.go
+++ b/internal/lsp/cmd/cmd_test.go
@@ -20,6 +20,10 @@
 )
 
 func TestMain(m *testing.M) {
+	if runtime.GOOS == "android" {
+		fmt.Fprintln(os.Stderr, "skipping test: cmd tests open too many files on android")
+		os.Exit(0)
+	}
 	testenv.ExitIfSmallMachine()
 	os.Exit(m.Run())
 }
diff --git a/internal/lsp/cmd/suggested_fix.go b/internal/lsp/cmd/suggested_fix.go
index 034cd62..19a53dc 100644
--- a/internal/lsp/cmd/suggested_fix.go
+++ b/internal/lsp/cmd/suggested_fix.go
@@ -65,6 +65,8 @@
 	if err := conn.diagnoseFiles(ctx, []span.URI{uri}); err != nil {
 		return err
 	}
+	conn.Client.filesMu.Lock()
+	defer conn.Client.filesMu.Unlock()
 
 	p := protocol.CodeActionParams{
 		TextDocument: protocol.TextDocumentIdentifier{
diff --git a/internal/lsp/diagnostics.go b/internal/lsp/diagnostics.go
index 2bcf640..cdbe16d 100644
--- a/internal/lsp/diagnostics.go
+++ b/internal/lsp/diagnostics.go
@@ -14,68 +14,76 @@
 	"golang.org/x/tools/internal/lsp/telemetry"
 	"golang.org/x/tools/internal/telemetry/log"
 	"golang.org/x/tools/internal/telemetry/trace"
+	"golang.org/x/tools/internal/xcontext"
 )
 
-func (s *Server) diagnoseSnapshot(ctx context.Context, snapshot source.Snapshot) {
-	ctx, done := trace.StartSpan(ctx, "lsp:background-worker")
-	defer done()
+func (s *Server) diagnoseDetached(snapshot source.Snapshot) {
+	ctx := snapshot.View().BackgroundContext()
+	ctx = xcontext.Detach(ctx)
 
-	wsPackages, err := snapshot.WorkspacePackages(ctx)
-	if err != nil {
-		log.Error(ctx, "diagnoseSnapshot: no workspace packages", err, telemetry.Directory.Of(snapshot.View().Folder))
-		return
-	}
-	for _, ph := range wsPackages {
-		go func(ph source.PackageHandle) {
-			reports, _, err := source.PackageDiagnostics(ctx, snapshot, ph, false, snapshot.View().Options().DisabledAnalyses)
-			if err != nil {
-				log.Error(ctx, "diagnoseSnapshot: no diagnostics", err, telemetry.Package.Of(ph.ID()))
-				return
-			}
-			s.publishReports(ctx, snapshot, reports, false)
-		}(ph)
-	}
-	// Run diagnostics on the go.mod file.
-	s.diagnoseModfile(snapshot)
+	s.diagnose(ctx, snapshot)
 }
 
-func (s *Server) diagnoseFile(snapshot source.Snapshot, fh source.FileHandle) {
+func (s *Server) diagnoseSnapshot(snapshot source.Snapshot) {
 	ctx := snapshot.View().BackgroundContext()
-	ctx, done := trace.StartSpan(ctx, "lsp:background-worker")
-	defer done()
 
-	ctx = telemetry.File.With(ctx, fh.Identity().URI)
-
-	reports, warningMsg, err := source.FileDiagnostics(ctx, snapshot, fh, true, snapshot.View().Options().DisabledAnalyses)
-	// Check the warning message first.
-	if warningMsg != "" {
-		s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
-			Type:    protocol.Info,
-			Message: warningMsg,
-		})
-	}
-	if err != nil {
-		if err != context.Canceled {
-			log.Error(ctx, "diagnoseFile: could not generate diagnostics", err)
-		}
-		return
-	}
-	s.publishReports(ctx, snapshot, reports, true)
+	s.diagnose(ctx, snapshot)
 }
 
-func (s *Server) diagnoseModfile(snapshot source.Snapshot) {
-	ctx := snapshot.View().BackgroundContext()
+// diagnose is a helper function for running diagnostics with a given context.
+// Do not call it directly.
+func (s *Server) diagnose(ctx context.Context, snapshot source.Snapshot) {
 	ctx, done := trace.StartSpan(ctx, "lsp:background-worker")
 	defer done()
 
-	reports, err := mod.Diagnostics(ctx, snapshot)
-	if err != nil {
-		if err != context.Canceled {
-			log.Error(ctx, "diagnoseModfile: could not generate diagnostics", err)
+	// Diagnose all of the packages in the workspace.
+	go func() {
+		wsPackages, err := snapshot.WorkspacePackages(ctx)
+		if err != nil {
+			log.Error(ctx, "diagnose: no workspace packages", err, telemetry.Directory.Of(snapshot.View().Folder))
+			return
 		}
-		return
-	}
-	s.publishReports(ctx, snapshot, reports, false)
+		for _, ph := range wsPackages {
+			go func(ph source.PackageHandle) {
+				// Only run analyses for packages with open files.
+				var withAnalyses bool
+				for _, fh := range ph.CompiledGoFiles() {
+					if s.session.IsOpen(fh.File().Identity().URI) {
+						withAnalyses = true
+					}
+				}
+				reports, msg, err := source.Diagnostics(ctx, snapshot, ph, withAnalyses)
+				// Check the warning message before the errors.
+				if msg != "" {
+					s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
+						Type:    protocol.Warning,
+						Message: msg,
+					})
+				}
+				if ctx.Err() != nil {
+					return
+				}
+				if err != nil {
+					log.Error(ctx, "diagnose: could not generate diagnostics for package", err, telemetry.Package.Of(ph.ID()))
+					return
+				}
+				s.publishReports(ctx, snapshot, reports, withAnalyses)
+			}(ph)
+		}
+	}()
+
+	// Diagnose the go.mod file.
+	go func() {
+		reports, err := mod.Diagnostics(ctx, snapshot)
+		if ctx.Err() != nil {
+			return
+		}
+		if err != nil {
+			log.Error(ctx, "diagnose: could not generate diagnostics for go.mod file", err)
+			return
+		}
+		s.publishReports(ctx, snapshot, reports, false)
+	}()
 }
 
 func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, reports map[source.FileIdentity][]source.Diagnostic, withAnalysis bool) {
diff --git a/internal/lsp/general.go b/internal/lsp/general.go
index b921590..065819d 100644
--- a/internal/lsp/general.go
+++ b/internal/lsp/general.go
@@ -17,7 +17,6 @@
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/telemetry/log"
-	"golang.org/x/tools/internal/xcontext"
 	errors "golang.org/x/xerrors"
 )
 
@@ -166,14 +165,12 @@
 
 	for _, folder := range folders {
 		uri := span.NewURI(folder.URI)
-		view, snapshot, err := s.addView(ctx, folder.Name, span.NewURI(folder.URI))
+		_, snapshot, err := s.addView(ctx, folder.Name, span.NewURI(folder.URI))
 		if err != nil {
 			viewErrors[uri] = err
 			continue
 		}
-		// Make sure that this does not get canceled.
-		ctx := xcontext.Detach(view.BackgroundContext())
-		go s.diagnoseSnapshot(ctx, snapshot)
+		go s.diagnoseDetached(snapshot)
 	}
 	if len(viewErrors) > 0 {
 		errMsg := fmt.Sprintf("Error loading workspace folders (expected %v, got %v)\n", len(folders), len(s.session.Views())-originalViews)
diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index 060cf58..07d44f7 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -88,16 +88,10 @@
 // TODO: Actually test the LSP diagnostics function in this test.
 func (r *runner) Diagnostics(t *testing.T, uri span.URI, want []source.Diagnostic) {
 	v := r.server.session.View(viewName)
-	fh, err := v.Snapshot().GetFile(uri)
+	_, got, err := source.FileDiagnostics(r.ctx, v.Snapshot(), uri)
 	if err != nil {
 		t.Fatal(err)
 	}
-	identity := fh.Identity()
-	results, _, err := source.FileDiagnostics(r.ctx, v.Snapshot(), fh, true, nil)
-	if err != nil {
-		t.Fatal(err)
-	}
-	got := results[identity]
 	// A special case to test that there are no diagnostics for a file.
 	if len(want) == 1 && want[0].Source == "no_diagnostics" {
 		if len(got) != 0 {
@@ -342,12 +336,7 @@
 		t.Fatal(err)
 	}
 	snapshot := view.Snapshot()
-	fh, err := snapshot.GetFile(spn.URI())
-	if err != nil {
-		t.Fatal(err)
-	}
-	fileID := fh.Identity()
-	diagnostics, _, err := source.FileDiagnostics(r.ctx, snapshot, fh, true, nil)
+	_, diagnostics, err := source.FileDiagnostics(r.ctx, snapshot, uri)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -357,7 +346,7 @@
 		},
 		Context: protocol.CodeActionContext{
 			Only:        []protocol.CodeActionKind{protocol.QuickFix},
-			Diagnostics: toProtocolDiagnostics(diagnostics[fileID]),
+			Diagnostics: toProtocolDiagnostics(diagnostics),
 		},
 	})
 	if err != nil {
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index d279012..fa3c77f 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -15,7 +15,6 @@
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/span"
-	errors "golang.org/x/xerrors"
 )
 
 // NewClientServer
@@ -293,27 +292,16 @@
 func (s *Server) NonstandardRequest(ctx context.Context, method string, params interface{}) (interface{}, error) {
 	paramMap := params.(map[string]interface{})
 	if method == "gopls/diagnoseFiles" {
-		files := paramMap["files"].([]interface{})
-		for _, file := range files {
+		for _, file := range paramMap["files"].([]interface{}) {
 			uri := span.URI(file.(string))
 			view, err := s.session.ViewOf(uri)
 			if err != nil {
 				return nil, err
 			}
-			snapshot := view.Snapshot()
-			fh, err := snapshot.GetFile(uri)
+			fileID, diagnostics, err := source.FileDiagnostics(ctx, view.Snapshot(), uri)
 			if err != nil {
 				return nil, err
 			}
-			reports, _, err := source.FileDiagnostics(ctx, view.Snapshot(), fh, true, nil)
-			if err != nil {
-				return nil, err
-			}
-			fileID := fh.Identity()
-			diagnostics, ok := reports[fileID]
-			if !ok {
-				return nil, errors.Errorf("no diagnostics for %s", uri)
-			}
 			if err := s.client.PublishDiagnostics(ctx, &protocol.PublishDiagnosticsParams{
 				URI:         protocol.NewURI(uri),
 				Diagnostics: toProtocolDiagnostics(diagnostics),
diff --git a/internal/lsp/source/diagnostics.go b/internal/lsp/source/diagnostics.go
index 8336706..fc8af76 100644
--- a/internal/lsp/source/diagnostics.go
+++ b/internal/lsp/source/diagnostics.go
@@ -40,44 +40,27 @@
 	Message string
 }
 
-func FileDiagnostics(ctx context.Context, snapshot Snapshot, fh FileHandle, withAnalysis bool, disabledAnalyses map[string]struct{}) (map[FileIdentity][]Diagnostic, string, error) {
-	if fh.Identity().Kind != Go {
-		return nil, "", errors.Errorf("unexpected file type: %q", fh.Identity().URI.Filename)
-	}
-	phs, err := snapshot.PackageHandles(ctx, fh)
-	if err != nil {
-		return nil, "", err
-	}
-	ph, err := WidestPackageHandle(phs)
-	if err != nil {
-		return nil, "", err
-	}
+func Diagnostics(ctx context.Context, snapshot Snapshot, ph PackageHandle, withAnalysis bool) (map[FileIdentity][]Diagnostic, string, error) {
 	// If we are missing dependencies, it may because the user's workspace is
 	// not correctly configured. Report errors, if possible.
-	var warningMsg string
+	var (
+		warn       bool
+		warningMsg string
+	)
 	if len(ph.MissingDependencies()) > 0 {
-		warningMsg, err = checkCommonErrors(ctx, snapshot.View())
-		if err != nil {
-			return nil, "", err
-		}
+		warn = true
 	}
-	reports, msg, err := PackageDiagnostics(ctx, snapshot, ph, withAnalysis, disabledAnalyses)
-	if warningMsg == "" {
-		warningMsg = msg
-	}
-	return reports, warningMsg, err
-}
-
-func PackageDiagnostics(ctx context.Context, snapshot Snapshot, ph PackageHandle, withAnalysis bool, disabledAnalyses map[string]struct{}) (map[FileIdentity][]Diagnostic, string, error) {
 	pkg, err := ph.Check(ctx)
 	if err != nil {
 		return nil, "", err
 	}
-	var warningMsg string
 	// If we have a package with a single file and errors about "undeclared" symbols,
 	// we may have an ad-hoc package with multiple files. Show a warning message.
 	// TODO(golang/go#36416): Remove this when golang.org/cl/202277 is merged.
 	if len(pkg.CompiledGoFiles()) == 1 && hasUndeclaredErrors(pkg) {
+		warn = true
+	}
+	if warn {
 		warningMsg, err = checkCommonErrors(ctx, snapshot.View())
 		if err != nil {
 			return nil, "", err
@@ -115,36 +98,41 @@
 	}
 	if !hadDiagnostics && withAnalysis {
 		// If we don't have any list, parse, or type errors, run analyses.
-		if err := analyses(ctx, snapshot, reports, ph, disabledAnalyses); err != nil {
+		if err := analyses(ctx, snapshot, reports, ph, snapshot.View().Options().DisabledAnalyses); err != nil {
 			// Exit early if the context has been canceled.
-			if err == context.Canceled {
-				return nil, warningMsg, err
+			if ctx.Err() != nil {
+				return nil, warningMsg, ctx.Err()
 			}
 			log.Error(ctx, "failed to run analyses", err, telemetry.Package.Of(ph.ID()))
 		}
 	}
-	// Updates to the diagnostics for this package may need to be propagated.
-	reverseDeps, err := snapshot.GetReverseDependencies(ctx, pkg.ID())
-	if err != nil {
-		if err == context.Canceled {
-			return nil, warningMsg, err
-		}
-		log.Error(ctx, "no reverse dependencies", err)
-		return reports, warningMsg, nil
-	}
-	for _, ph := range reverseDeps {
-		pkg, err := ph.Check(ctx)
-		if err != nil {
-			return nil, warningMsg, err
-		}
-		for _, fh := range pkg.CompiledGoFiles() {
-			clearReports(snapshot, reports, fh.File().Identity())
-		}
-		diagnostics(ctx, snapshot, reports, pkg)
-	}
 	return reports, warningMsg, nil
 }
 
+func FileDiagnostics(ctx context.Context, snapshot Snapshot, uri span.URI) (FileIdentity, []Diagnostic, error) {
+	fh, err := snapshot.GetFile(uri)
+	if err != nil {
+		return FileIdentity{}, nil, err
+	}
+	phs, err := snapshot.PackageHandles(ctx, fh)
+	if err != nil {
+		return FileIdentity{}, nil, err
+	}
+	ph, err := NarrowestPackageHandle(phs)
+	if err != nil {
+		return FileIdentity{}, nil, err
+	}
+	reports, _, err := Diagnostics(ctx, snapshot, ph, true)
+	if err != nil {
+		return FileIdentity{}, nil, err
+	}
+	diagnostics, ok := reports[fh.Identity()]
+	if !ok {
+		return FileIdentity{}, nil, errors.Errorf("no diagnostics for %s", uri)
+	}
+	return fh.Identity(), diagnostics, nil
+}
+
 type diagnosticSet struct {
 	listErrors, parseErrors, typeErrors []*Diagnostic
 }
diff --git a/internal/lsp/source/errors.go b/internal/lsp/source/errors.go
index 5d2d534..4eed4c0 100644
--- a/internal/lsp/source/errors.go
+++ b/internal/lsp/source/errors.go
@@ -35,9 +35,10 @@
 	// TODO(rstambler): Get the values for GOPATH and GOMOD from
 	// the view, once it's possible to do so: golang.org/cl/214417.
 	gopath := os.Getenv("GOPATH")
+	folder := v.Folder().Filename()
 
-	// Invoke `go env GOMOD` inside of the directory of the file.
-	b, err := InvokeGo(ctx, v.Folder().Filename(), v.Config(ctx).Env, "env", "GOMOD")
+	// Invoke `go env GOMOD` inside of the directory of the view.
+	b, err := InvokeGo(ctx, folder, v.Config(ctx).Env, "env", "GOMOD")
 	if err != nil {
 		return "", err
 	}
@@ -49,7 +50,6 @@
 
 	// Not inside of a module.
 	inAModule := modFile != ""
-	folder := v.Folder().Filename()
 
 	// The user may have a multiple directories in their GOPATH.
 	var inGopath bool
diff --git a/internal/lsp/source/source_test.go b/internal/lsp/source/source_test.go
index 1c13d2f..bc95fc6 100644
--- a/internal/lsp/source/source_test.go
+++ b/internal/lsp/source/source_test.go
@@ -80,16 +80,11 @@
 
 func (r *runner) Diagnostics(t *testing.T, uri span.URI, want []source.Diagnostic) {
 	snapshot := r.view.Snapshot()
-	fh, err := snapshot.GetFile(uri)
+
+	fileID, got, err := source.FileDiagnostics(r.ctx, snapshot, uri)
 	if err != nil {
 		t.Fatal(err)
 	}
-	fileID := fh.Identity()
-	results, _, err := source.FileDiagnostics(r.ctx, snapshot, fh, true, nil)
-	if err != nil {
-		t.Fatal(err)
-	}
-	got := results[fileID]
 	// A special case to test that there are no diagnostics for a file.
 	if len(want) == 1 && want[0].Source == "no_diagnostics" {
 		if len(got) != 0 {
diff --git a/internal/lsp/source/util.go b/internal/lsp/source/util.go
index 9e40ca0..9d729ec 100644
--- a/internal/lsp/source/util.go
+++ b/internal/lsp/source/util.go
@@ -141,17 +141,17 @@
 	}
 }
 
-func IsGenerated(ctx context.Context, view View, uri span.URI) bool {
-	fh, err := view.Snapshot().GetFile(uri)
+func IsGenerated(ctx context.Context, snapshot Snapshot, uri span.URI) bool {
+	fh, err := snapshot.GetFile(uri)
 	if err != nil {
 		return false
 	}
-	ph := view.Session().Cache().ParseGoHandle(fh, ParseHeader)
+	ph := snapshot.View().Session().Cache().ParseGoHandle(fh, ParseHeader)
 	parsed, _, _, err := ph.Parse(ctx)
 	if err != nil {
 		return false
 	}
-	tok := view.Session().Cache().FileSet().File(parsed.Pos())
+	tok := snapshot.View().Session().Cache().FileSet().File(parsed.Pos())
 	if tok == nil {
 		return false
 	}
diff --git a/internal/lsp/text_synchronization.go b/internal/lsp/text_synchronization.go
index 93e588d..96fe881 100644
--- a/internal/lsp/text_synchronization.go
+++ b/internal/lsp/text_synchronization.go
@@ -18,32 +18,14 @@
 )
 
 func (s *Server) didOpen(ctx context.Context, params *protocol.DidOpenTextDocumentParams) error {
-	// Confirm that the file's language ID is related to Go.
-	uri := span.NewURI(params.TextDocument.URI)
-	snapshots, err := s.session.DidModifyFile(ctx, source.FileModification{
-		URI:        uri,
+	_, err := s.didModifyFile(ctx, source.FileModification{
+		URI:        span.NewURI(params.TextDocument.URI),
 		Action:     source.Open,
 		Version:    params.TextDocument.Version,
 		Text:       []byte(params.TextDocument.Text),
 		LanguageID: params.TextDocument.LanguageID,
 	})
-	if err != nil {
-		return err
-	}
-	snapshot, _, err := snapshotOf(s.session, uri, snapshots)
-	if err != nil {
-		return err
-	}
-	fh, err := snapshot.GetFile(uri)
-	if err != nil {
-		return err
-	}
-	// Only run diagnostics on file open for Go files.
-	switch fh.Identity().Kind {
-	case source.Go:
-		go s.diagnoseFile(snapshot, fh)
-	}
-	return nil
+	return err
 }
 
 func (s *Server) didChange(ctx context.Context, params *protocol.DidChangeTextDocumentParams) error {
@@ -52,38 +34,25 @@
 	if err != nil {
 		return err
 	}
-	snapshots, err := s.session.DidModifyFile(ctx, source.FileModification{
+	c := source.FileModification{
 		URI:     uri,
 		Action:  source.Change,
 		Version: params.TextDocument.Version,
 		Text:    text,
-	})
-	if err != nil {
-		return err
 	}
-	snapshot, view, err := snapshotOf(s.session, uri, snapshots)
+	snapshot, err := s.didModifyFile(ctx, c)
 	if err != nil {
 		return err
 	}
 	// Ideally, we should be able to specify that a generated file should be opened as read-only.
 	// Tell the user that they should not be editing a generated file.
-	if s.wasFirstChange(uri) && source.IsGenerated(ctx, view, uri) {
-		s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
+	if s.wasFirstChange(uri) && source.IsGenerated(ctx, snapshot, uri) {
+		if err := s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
 			Message: fmt.Sprintf("Do not edit this file! %s is a generated file.", uri.Filename()),
 			Type:    protocol.Warning,
-		})
-	}
-	fh, err := snapshot.GetFile(uri)
-	if err != nil {
-		return err
-	}
-	// Run diagnostics for Go files and for mod files.
-	switch fh.Identity().Kind {
-	case source.Go:
-		go s.diagnoseFile(snapshot, fh)
-	case source.Mod:
-		// TODO(rstambler): Modifying the go.mod file should trigger workspace-level diagnostics.
-		go s.diagnoseModfile(snapshot)
+		}); err != nil {
+			return err
+		}
 	}
 	return nil
 }
@@ -115,8 +84,8 @@
 		views[snapshot.View()] = snapshot
 	}
 	// Diagnose all resulting snapshots.
-	for view, snapshot := range views {
-		go s.diagnoseSnapshot(view.BackgroundContext(), snapshot)
+	for _, snapshot := range views {
+		go s.diagnoseSnapshot(snapshot)
 	}
 	return nil
 }
@@ -130,20 +99,12 @@
 	if params.Text != nil {
 		c.Text = []byte(*params.Text)
 	}
-	snapshots, err := s.session.DidModifyFile(ctx, c)
-	if err != nil {
-		return err
-	}
-	snapshot, _, err := snapshotOf(s.session, c.URI, snapshots)
-	if err != nil {
-		return err
-	}
-	go s.diagnoseModfile(snapshot)
+	_, err := s.didModifyFile(ctx, c)
 	return err
 }
 
 func (s *Server) didClose(ctx context.Context, params *protocol.DidCloseTextDocumentParams) error {
-	_, err := s.session.DidModifyFile(ctx, source.FileModification{
+	_, err := s.didModifyFile(ctx, source.FileModification{
 		URI:     span.NewURI(params.TextDocument.URI),
 		Action:  source.Close,
 		Version: -1,
@@ -152,6 +113,33 @@
 	return err
 }
 
+func (s *Server) didModifyFile(ctx context.Context, c source.FileModification) (source.Snapshot, error) {
+	snapshots, err := s.session.DidModifyFile(ctx, c)
+	if err != nil {
+		return nil, err
+	}
+	snapshot, _, err := snapshotOf(s.session, c.URI, snapshots)
+	if err != nil {
+		return nil, err
+	}
+	switch c.Action {
+	case source.Save:
+		// If we're saving a go.mod file, all of the metadata has been invalidated,
+		// so we need to run diagnostics and make sure that they cannot be canceled.
+		fh, err := snapshot.GetFile(c.URI)
+		if err != nil {
+			return nil, err
+		}
+		if fh.Identity().Kind == source.Mod {
+			go s.diagnoseDetached(snapshot)
+		}
+	default:
+		go s.diagnoseSnapshot(snapshot)
+	}
+
+	return snapshot, nil
+}
+
 // snapshotOf returns the snapshot corresponding to the view for the given file URI.
 func snapshotOf(session source.Session, uri span.URI, snapshots []source.Snapshot) (source.Snapshot, source.View, error) {
 	view, err := session.ViewOf(uri)
@@ -184,7 +172,6 @@
 	if len(changes) == 1 && changes[0].Range == nil && changes[0].RangeLength == 0 {
 		return []byte(changes[0].Text), nil
 	}
-
 	return s.applyIncrementalChanges(ctx, uri, changes)
 }
 
diff --git a/internal/lsp/workspace.go b/internal/lsp/workspace.go
index a607d9e..2a5d2b0 100644
--- a/internal/lsp/workspace.go
+++ b/internal/lsp/workspace.go
@@ -10,7 +10,6 @@
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/span"
-	"golang.org/x/tools/internal/xcontext"
 	errors "golang.org/x/xerrors"
 )
 
@@ -53,9 +52,7 @@
 		if err != nil {
 			return err
 		}
-		// Make sure that this does not get canceled.
-		ctx := xcontext.Detach(view.BackgroundContext())
-		go s.diagnoseSnapshot(ctx, view.Snapshot())
+		go s.diagnoseDetached(view.Snapshot())
 	}
 	return nil
 }