gopls/internal/server: simplify DiagnoseFiles to avoid a race

Simplify the DiagnoseFiles command handler to just diagnose all
snapshots for requested files, rather than implement ad-hoc diagnostic
logic. This may be somewhat slower, but this is only used for command
line commands and is therefore performance is not critical.

Also adjust updateDiagnostics to not overwrite final diagnostics with
non-final diagnostics.

These two changes together avoid the race encountered in
golang/go#64765: DiagnoseFiles does not return until diagnostics are
finalized, and these final diagnostics are not overwritten.

Fixes golang/go#64765

Change-Id: I54ef7309487a9803a8bbd45ab2a8de4dbf30c460
Reviewed-on: https://go-review.googlesource.com/c/tools/+/556475
Reviewed-by: Alan Donovan <[email protected]>
Run-TryBot: Robert Findley <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>

Author: findleyr

gopls/internal/lsp/cache/pkg.go

@@ -5,7 +5,6 @@
 package cache
 
 import (
-	"context"
 	"fmt"
 	"go/ast"
 	"go/scanner"
@@ -172,19 +171,3 @@ func (p *Package) GetParseErrors() []scanner.ErrorList {
 func (p *Package) GetTypeErrors() []types.Error {
 	return p.pkg.typeErrors
 }
-
-func (p *Package) DiagnosticsForFile(ctx context.Context, uri protocol.DocumentURI) ([]*Diagnostic, error) {
-	var diags []*Diagnostic
-	for _, diag := range p.loadDiagnostics {
-		if diag.URI == uri {
-			diags = append(diags, diag)
-		}
-	}
-	for _, diag := range p.pkg.diagnostics {
-		if diag.URI == uri {
-			diags = append(diags, diag)
-		}
-	}
-
-	return diags, nil
-}

gopls/internal/server/command.go

@@ -34,7 +34,6 @@ import (
 	"golang.org/x/tools/gopls/internal/settings"
 	"golang.org/x/tools/gopls/internal/telemetry"
 	"golang.org/x/tools/gopls/internal/util/bug"
-	"golang.org/x/tools/gopls/internal/util/maps"
 	"golang.org/x/tools/gopls/internal/vulncheck"
 	"golang.org/x/tools/gopls/internal/vulncheck/scan"
 	"golang.org/x/tools/internal/diff"
@@ -1301,47 +1300,31 @@ func (c *commandHandler) DiagnoseFiles(ctx context.Context, args command.Diagnos
 		ctx, done := event.Start(ctx, "lsp.server.DiagnoseFiles")
 		defer done()
 
-		// TODO(adonovan): opt: parallelize the loop,
-		// grouping file URIs by package and making a
-		// single call to source.Analyze.
+		snapshots := make(map[*cache.Snapshot]bool)
 		for _, uri := range args.Files {
 			fh, snapshot, release, err := c.s.fileOf(ctx, uri)
 			if err != nil {
 				return err
 			}
-			defer release()
-			if snapshot.FileKind(fh) != file.Go {
+			if snapshots[snapshot] || snapshot.FileKind(fh) != file.Go {
+				release()
 				continue
 			}
-			pkg, _, err := source.NarrowestPackageForFile(ctx, snapshot, uri)
-			if err != nil {
-				return err
-			}
-			pkgDiags, err := pkg.DiagnosticsForFile(ctx, uri)
-			if err != nil {
-				return err
-			}
-			adiags, err := source.Analyze(ctx, snapshot, map[source.PackageID]unit{pkg.Metadata().ID: {}}, nil /* progress tracker */)
-			if err != nil {
-				return err
-			}
+			defer release()
+			snapshots[snapshot] = true
+		}
 
-			// combine load/parse/type + analysis diagnostics
-			var td, ad []*cache.Diagnostic
-			combineDiagnostics(pkgDiags, adiags[uri], &td, &ad)
-			diags := append(td, ad...)
-			byURI := func(d *cache.Diagnostic) protocol.DocumentURI { return d.URI }
-			c.s.updateDiagnostics(ctx, c.s.session.Views(), snapshot, maps.Group(diags, byURI), false)
-			diagnostics := append(td, ad...)
-
-			if err := c.s.client.PublishDiagnostics(ctx, &protocol.PublishDiagnosticsParams{
-				URI:         fh.URI(),
-				Version:     fh.Version(),
-				Diagnostics: toProtocolDiagnostics(diagnostics),
-			}); err != nil {
-				return err
-			}
+		var wg sync.WaitGroup
+		for snapshot := range snapshots {
+			snapshot := snapshot
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				c.s.diagnoseSnapshot(snapshot, nil, 0)
+			}()
 		}
+		wg.Wait()
+
 		return nil
 	})
 }

gopls/internal/server/diagnostics.go

@@ -674,7 +674,15 @@ func (s *server) updateDiagnostics(ctx context.Context, allViews []*cache.View,
 	// views.
 	updateAndPublish := func(uri protocol.DocumentURI, f *fileDiagnostics, diags []*cache.Diagnostic) error {
 		current, ok := f.byView[snapshot.View()]
-		if !ok || current.snapshot <= snapshot.SequenceID() {
+		// Update the stored diagnostics if:
+		//  1. we've never seen diagnostics for this view,
+		//  2. diagnostics are for an older snapshot, or
+		//  3. we're overwriting with final diagnostics
+		//
+		// In other words, we shouldn't overwrite existing diagnostics for a
+		// snapshot with non-final diagnostics. This avoids the race described at
+		// https://github.com/golang/go/issues/64765#issuecomment-1890144575.
+		if !ok || current.snapshot < snapshot.SequenceID() || (current.snapshot == snapshot.SequenceID() && final) {
 			fh, err := snapshot.ReadFile(ctx, uri)
 			if err != nil {
 				return err