internal/lsp: aggregate diagnostics before publishing them

This change aggregates all the diagnostics before publishing them to the client. This also updates the lsp_test to use lsp.Diagnose vs source.FileDiagnostics, unless lsp.Diagnose doesn't return any results.

Change-Id: I39959b50af3f65fce5d8e15fb138ccc019ce8073
Reviewed-on: https://go-review.googlesource.com/c/tools/+/217338
Run-TryBot: Rohan Challa <rohan@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/diagnostics.go

@@ -7,6 +7,7 @@ package lsp
 import (
 	"context"
 	"strings"
+	"sync"
 
 	"golang.org/x/tools/internal/lsp/mod"
 	"golang.org/x/tools/internal/lsp/protocol"
@@ -17,33 +18,44 @@ import (
 	"golang.org/x/tools/internal/xcontext"
 )
 
+type diagnosticKey struct {
+	id           source.FileIdentity
+	withAnalysis bool
+}
+
 func (s *Server) diagnoseDetached(snapshot source.Snapshot) {
 	ctx := snapshot.View().BackgroundContext()
 	ctx = xcontext.Detach(ctx)
 
-	s.diagnose(ctx, snapshot)
+	reports := s.diagnose(ctx, snapshot, false)
+	s.publishReports(ctx, snapshot, reports)
 }
 
 func (s *Server) diagnoseSnapshot(snapshot source.Snapshot) {
 	ctx := snapshot.View().BackgroundContext()
 
-	s.diagnose(ctx, snapshot)
+	reports := s.diagnose(ctx, snapshot, false)
+	s.publishReports(ctx, snapshot, reports)
 }
 
 // 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) {
+func (s *Server) diagnose(ctx context.Context, snapshot source.Snapshot, alwaysAnalyze bool) map[diagnosticKey][]source.Diagnostic {
 	ctx, done := trace.StartSpan(ctx, "lsp:background-worker")
 	defer done()
 
+	allReports := make(map[diagnosticKey][]source.Diagnostic)
+	var reportsMu sync.Mutex
+	var wg sync.WaitGroup
+
 	// Diagnose the go.mod file.
 	reports, missingModules, err := mod.Diagnostics(ctx, snapshot)
 	if ctx.Err() != nil {
-		return
+		return nil
 	}
 	if err != nil {
 		log.Error(ctx, "diagnose: could not generate diagnostics for go.mod file", err)
-		return
+		return nil
 	}
 	// Ensure that the reports returned from mod.Diagnostics are only related to the
 	// go.mod file for the module.
@@ -51,55 +63,68 @@ func (s *Server) diagnose(ctx context.Context, snapshot source.Snapshot) {
 		panic("unexpected reports from mod.Diagnostics")
 	}
 	modURI, _ := snapshot.View().ModFiles()
-	for fi := range reports {
-		if fi.URI != modURI {
+	for id, diags := range reports {
+		if id.URI != modURI {
 			panic("unexpected reports from mod.Diagnostics")
 		}
+		key := diagnosticKey{
+			id: id,
+		}
+		allReports[key] = diags
 	}
-	s.publishReports(ctx, snapshot, reports, false)
 
 	// Diagnose all of the packages in the workspace.
-	go func() {
-		wsPackages, err := snapshot.WorkspacePackages(ctx)
-		if ctx.Err() != nil {
-			return
-		}
-		if err != nil {
-			log.Error(ctx, "diagnose: no workspace packages", err, telemetry.Snapshot.Of(snapshot.ID()), telemetry.Directory.Of(snapshot.View().Folder))
-			return
-		}
-		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
-					}
+	wsPackages, err := snapshot.WorkspacePackages(ctx)
+	if ctx.Err() != nil {
+		return nil
+	}
+	if err != nil {
+		log.Error(ctx, "diagnose: no workspace packages", err, telemetry.Snapshot.Of(snapshot.ID()), telemetry.Directory.Of(snapshot.View().Folder))
+		return nil
+	}
+	for _, ph := range wsPackages {
+		wg.Add(1)
+		go func(ph source.PackageHandle) {
+			defer wg.Done()
+			// Only run analyses for packages with open files.
+			withAnalyses := alwaysAnalyze
+			for _, fh := range ph.CompiledGoFiles() {
+				if s.session.IsOpen(fh.File().Identity().URI) {
+					withAnalyses = true
 				}
-				reports, warn, err := source.Diagnostics(ctx, snapshot, ph, missingModules, withAnalyses)
-				// Check if might want to warn the user about their build configuration.
-				if warn && !snapshot.View().ValidBuildConfiguration() {
-					s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
-						Type: protocol.Warning,
-						// TODO(rstambler): We should really be able to point to a link on the website.
-						Message: `You are neither in a module nor in your GOPATH. Please see https://github.com/golang/go/wiki/Modules for information on how to set up your Go project.`,
-					})
+			}
+			reports, warn, err := source.Diagnostics(ctx, snapshot, ph, missingModules, withAnalyses)
+			// Check if might want to warn the user about their build configuration.
+			if warn && !snapshot.View().ValidBuildConfiguration() {
+				s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
+					Type: protocol.Warning,
+					// TODO(rstambler): We should really be able to point to a link on the website.
+					Message: `You are neither in a module nor in your GOPATH. Please see https://github.com/golang/go/wiki/Modules for information on how to set up your Go project.`,
+				})
+			}
+			if ctx.Err() != nil {
+				return
+			}
+			if err != nil {
+				log.Error(ctx, "diagnose: could not generate diagnostics for package", err, telemetry.Snapshot.Of(snapshot.ID()), telemetry.Package.Of(ph.ID()))
+				return
+			}
+			reportsMu.Lock()
+			for id, diags := range reports {
+				key := diagnosticKey{
+					id:           id,
+					withAnalysis: withAnalyses,
 				}
-				if ctx.Err() != nil {
-					return
-				}
-				if err != nil {
-					log.Error(ctx, "diagnose: could not generate diagnostics for package", err, telemetry.Snapshot.Of(snapshot.ID()), telemetry.Package.Of(ph.ID()))
-					return
-				}
-				s.publishReports(ctx, snapshot, reports, withAnalyses)
-			}(ph)
-		}
-	}()
+				allReports[key] = diags
+			}
+			reportsMu.Unlock()
+		}(ph)
+	}
+	wg.Wait()
+	return allReports
 }
 
-func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, reports map[source.FileIdentity][]source.Diagnostic, withAnalysis bool) {
+func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, reports map[diagnosticKey][]source.Diagnostic) {
 	// Check for context cancellation before publishing diagnostics.
 	if ctx.Err() != nil {
 		return
@@ -108,7 +133,7 @@ func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, r
 	s.deliveredMu.Lock()
 	defer s.deliveredMu.Unlock()
 
-	for fileID, diagnostics := range reports {
+	for key, diagnostics := range reports {
 		// Don't deliver diagnostics if the context has already been canceled.
 		if ctx.Err() != nil {
 			break
@@ -117,15 +142,15 @@ func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, r
 		// Pre-sort diagnostics to avoid extra work when we compare them.
 		source.SortDiagnostics(diagnostics)
 		toSend := sentDiagnostics{
-			version:      fileID.Version,
-			identifier:   fileID.Identifier,
+			version:      key.id.Version,
+			identifier:   key.id.Identifier,
 			sorted:       diagnostics,
-			withAnalysis: withAnalysis,
+			withAnalysis: key.withAnalysis,
 			snapshotID:   snapshot.ID(),
 		}
 
 		// We use the zero values if this is an unknown file.
-		delivered := s.delivered[fileID.URI]
+		delivered := s.delivered[key.id.URI]
 
 		// Snapshot IDs are always increasing, so we use them instead of file
 		// versions to create the correct order for diagnostics.
@@ -140,7 +165,7 @@ func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, r
 		// Check if we should reuse the cached diagnostics.
 		if equalDiagnostics(delivered.sorted, diagnostics) {
 			// Make sure to update the delivered map.
-			s.delivered[fileID.URI] = toSend
+			s.delivered[key.id.URI] = toSend
 			continue
 		}
 
@@ -154,8 +179,8 @@ func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, r
 
 		if err := s.client.PublishDiagnostics(ctx, &protocol.PublishDiagnosticsParams{
 			Diagnostics: toProtocolDiagnostics(diagnostics),
-			URI:         protocol.NewURI(fileID.URI),
-			Version:     fileID.Version,
+			URI:         protocol.NewURI(key.id.URI),
+			Version:     key.id.Version,
 		}); err != nil {
 			if ctx.Err() == nil {
 				log.Error(ctx, "publishReports: failed to deliver diagnostic", err, telemetry.File)
@@ -163,7 +188,7 @@ func (s *Server) publishReports(ctx context.Context, snapshot source.Snapshot, r
 			continue
 		}
 		// Update the delivered map.
-		s.delivered[fileID.URI] = toSend
+		s.delivered[key.id.URI] = toSend
 	}
 }
 

internal/lsp/lsp_test.go

@@ -39,9 +39,10 @@ func TestLSP(t *testing.T) {
 }
 
 type runner struct {
-	server *Server
-	data   *tests.Data
-	ctx    context.Context
+	server      *Server
+	data        *tests.Data
+	diagnostics map[span.URI][]source.Diagnostic
+	ctx         context.Context
 }
 
 const viewName = "lsp_test"
@@ -87,13 +88,18 @@ func testLSP(t *testing.T, exporter packagestest.Exporter) {
 	tests.Run(t, r, data)
 }
 
-// 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)
-	_, got, err := source.FileDiagnostics(r.ctx, v.Snapshot(), uri)
-	if err != nil {
-		t.Fatal(err)
+	// Get the diagnostics for this view if we have not done it before.
+	if r.diagnostics == nil {
+		r.diagnostics = make(map[span.URI][]source.Diagnostic)
+		v := r.server.session.View(viewName)
+		// Always run diagnostics with analysis.
+		reports := r.server.diagnose(r.ctx, v.Snapshot(), true)
+		for key, diags := range reports {
+			r.diagnostics[key.id.URI] = diags
+		}
 	}
+	got := r.diagnostics[uri]
 	// 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 {